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

AI summary

Overview

Research area: Generative modeling of neural network weights (weight-space learning), combining Flow Matching with neural network symmetry canonicalization.

Technical level: Advanced. The paper assumes familiarity with flow matching, permutation symmetries in weight space, canonicalization/model alignment, and the standard vision and language architectures it targets.

Scope: The paper introduces DeepWeightFlow, a Flow Matching model that generates complete, high-accuracy weight sets directly in weight space for MLPs, ResNets, ViTs, and BERT, and evaluates it on vision, tabular, and language tasks.

What This Paper Is About

Training a neural network is slow, and generative models that could simply sample trained weights would bypass much of that cost. Existing weight generators either produce only partial weight sets for large models, need fine-tuning after generation, or take hours per sample, and they struggle with the permutation symmetries that make many distinct weight configurations functionally identical. This paper presents DeepWeightFlow, a Flow Matching model that generates complete weight sets for full networks without fine-tuning, scales to roughly 100M parameters, and generates diverse ensembles in minutes.

Key Contributions

1. A complete-weight Flow Matching generator. DeepWeightFlow generates complete neural network weight sets, unconditioned by dataset characteristics, task descriptions, or architectural specifications. It requires no auxiliary autoencoder for dimensionality reduction and can scale to high-dimensional weight spaces using PCA instead.

2. Demonstrated scaling across architectures and modalities. The method generates weights for networks with 𝒪(100M) parameters across MLP, ResNet, ViT, and BERT, producing models that perform well on vision, tabular, and natural language tasks without fine-tuning.

3. An empirical account of symmetry's role. The authors show that canonicalizing the training data (via Git Re-Basin or TransFusion) helps when generating very high-dimensional weights, but provides no additional benefit when the weight-space dimension is moderate — a capacity-dependent effect they map out systematically.

4. Large efficiency gains over diffusion-based weight generators. Using a simple MLP implementation with no equivariant architecture, DeepWeightFlow trains in up to 𝒪(10) minutes and samples in seconds, versus the hours required by diffusion alternatives such as RPG, P-diff, and D2NWG.

Main Findings

• Complete generation matches the training set. On MNIST MLPs, DeepWeightFlow with Git Re-Basin reproduced an original accuracy of 96.32 ± 0.20 with 96.17 ± 0.31 generated; without Git Re-Basin the generated result was 96.19 ± 0.27. For ResNet-18 on CIFAR-10, the original was 94.45 ± 0.14 and the complete generated result 93.55 ± 0.13 with Git Re-Basin and 93.47 ± 0.20 without. For ResNet-18 on STL-10, generation reached 62.46 ± 0.79 with Re-Basin and 62.50 ± 0.66 without, against an original of 62.30 ± 0.77.

• Competitive against state-of-the-art weight generators. Listed comparisons include RPG (95.3 original, 95.1 generated), D2NWG (94.56 original, 94.57 ± 0.0 generated), SANE (92.14 ± 0.12 original, 68.6 ± 1.2 generated), N_M unconditioned (94.54 original, 94.36 generated), P-diff best network (94.54 original, 94.36 generated), and FLoWN best network (94.54 original, 94.36 generated) on ResNet-18 CIFAR-10. On ViT-Small-192 CIFAR-10, DeepWeightFlow reached 83.07 ± 0.42 generated (with TransFusion) versus 83.30 ± 0.29 original, and 82.58 ± 0.07 without TransFusion; P-diff best was reported at 73.6 generated on ViT-mini and RPG at 98.9 on ViT-Base.

• BERT generation works on a regression task. On Yelp Review with BERT-118M, Spearman correlation was 0.7902 ± 0.061 original, 0.7909 ± 0.005 generated with Re-Basin, and 0.7884 ± 0.012 without, at a flow hidden dimension of 1024; at 768 the corresponding figures were 0.7894 ± 0.006 and 0.7892 ± 0.015.

• Canonicalization helps at low flow-model capacity. In Table 5 the gap is largest when the flow hidden dimension is small — for example MNIST MLP at 𝑑_h = 64 (57.80 ± 9.85 with Re-Basin versus 25.54 ± 12.90 without), ViT-Small-192 on CIFAR-10 at 64 (43.13 ± 30.28 versus 12.67 ± 7.11), and ResNet-18 at 64 (29.92 ± 19.79 versus 21.93 ± 19.86). As capacity grows, canonicalized and non-canonicalized models converge to similar performance.

• Source distribution choice matters. Gaussian noise consistently outperformed alternatives such as Kaiming initialization as the source distribution, with best results when the source standard deviation matches or slightly undershoots that of the target weight distribution. The sensitivity is most pronounced in smaller flow models.

• Robust to initialization schemes. A single flow model trained on a heterogeneous Iris MLP collection of 100 models (20 seeds × 5 initialization types: Kaiming, Xavier, Kaiming weights with zero biases, normal, and uniform) generated novel weights achieving high test accuracy. All other experiments used Kaiming initialization with varied seeds.

• Generated models transfer well. ResNet-18 classifiers generated on CIFAR-10 (using PCA, with 5 models generated and BN statistics recalibrated on a small CIFAR-10 subset) were evaluated on STL-10 and SVHN. At epoch 0, generated models reached 48.32 ± 0.34 and 11.57 ± 0.49, essentially matching pretrained models (48.31 ± 0.17 and 11.51 ± 0.31) and far exceeding FLoWN generated models (35.16 ± 1.24 and 17.99 ± 0.82). At 5 epochs, generated models reached 84.63 ± 0.17 on STL-10 and 95.85 ± 0.09 on SVHN. Against SANE on SmallCNN, DeepWeightFlow generated models reached 62.62 ± 0.46 at 25 epochs versus a pretrained 62.14 ± 0.84 and SANE_SUB 49.8 ± 0.6.

• Generated networks are diverse, not memorized. Using maximum IoU between wrong-prediction sets, the authors show (with 500 generated networks in Figure 2) that generated models overlap with originals perturbed by 𝒩(0, 0.01) noise in the left panels, but the middle panels show generated networks are concretely different from originals with that noise added. Training sets consist of 100 terminal networks from unique seeds, not checkpoints from a single training run.

• Batch normalization needs recalibration. Directly transferring running statistics from a reference model yields suboptimal performance; DeepWeightFlow recomputes BN statistics on the training dataset for each generated weight set. Layer normalization is permutation invariant and does not need recalibration.

• Scaling uses PCA. Incremental PCA is used for weight-space dimensions of 𝒪(10M) and Dual PCA for 𝒪(100M), with inverse PCA during generation. The authors estimate 𝒪(1B)-parameter models may be feasible and leave that as future work.

• Efficiency. Training takes up to 𝒪(10) minutes for most architectures up to 𝒪(100M) parameters, versus several hours for RPG; generation takes seconds versus minutes or hours for RPG, P-Diff, or D2NWG. Ensembles of hundreds of networks can be generated in minutes.

Methodology in Plain English

The researchers first build a training set of weights by fully training many neural networks from random, independent initializations on a target task (not checkpoints from one training run), which keeps the set diverse. Optionally, they canonicalize these weights — Git Re-Basin for MLPs and ResNets (run for 100 iterations) and TransFusion for ViTs (run for 10 iterations, since it uses spectral decomposition and is slower) — to collapse permutation-equivalent weights into one representative per orbit.

They then train a time-conditioned MLP that predicts a velocity field carrying Gaussian noise to the weight distribution. Along a straight-line path between a noise sample and a trained weight vector, the model learns the constant direction of travel, conditioning on a learned embedding of the interpolation time. Sampling integrates this learned field from Gaussian noise using a fourth-order Runge-Kutta method. For larger models, incremental or Dual PCA compresses flattened weight vectors before training and inverse-transforms after generation to stay within GPU memory. After generation, batch normalization running statistics are recomputed on training data so the sampled weights are actually accurate.

Why This Matters

Research impact. The paper reframes weight generation as direct sampling in weight space rather than a latent-variable problem, and it is the first to systematically test how canonicalization interacts with flow-model capacity. It also extends canonicalization to transformers via TransFusion in the generative weight-model context, and demonstrates PCA as a lighter-weight alternative to autoencoders for scaling weight generation.

Real-world applications (drawn from the motivations the paper states):

• Model editing and targeted modification of trained networks.
• Accelerating transfer learning by supplying ready-to-finetune starting weights.
• Uncertainty quantification through ensembles of hundreds of diverse networks.
• Neural architecture search, where many candidate weight sets must be evaluated quickly.

Industry relevance. Training large models is a dominant compute cost, and a method that produces useful weights in minutes without post-hoc fine-tuning could reduce that cost, support ensemble-based robustness, and address privacy concerns tied to training-data reconstruction by avoiding repeated exposure to raw data during training.

Future Directions

• Extending DeepWeightFlow to models of 𝒪(1B) parameters, which the authors estimate may be possible with Dual PCA and explicitly leave as future work.
• Understanding and closing the residual gap to the original training set at low flow-model capacity, where canonicalization currently makes the largest difference.
• Testing whether conditioning on dataset characteristics or task descriptions (which DeepWeightFlow deliberately avoids) could extend the method to broader multi-task settings; the paper only probes class-conditioning in Subsection K.2.
• Making the generated training datasets publicly available, which the paper states will happen in the future, alongside the released code at the project's GitHub repository.

Target Audience

Researchers and practitioners in generative modeling, weight-space learning, and model merging who want complete weight generation without fine-tuning. It is most useful to those already comfortable with flow matching, permutation symmetries, and canonicalization methods such as Git Re-Basin and TransFusion, though the results tables and efficiency claims are readable by anyone evaluating whether weight generation is a practical alternative to training.

Authors’ 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.