DeepWeightFlow: Symmetry-Aware Weight Generation

• DeepWeightFlow is a flow matching model that generates neural network weights by addressing high-dimensional symmetries through canonicalization techniques.
• It employs direct vector field regression in weight space with linear interpolation paths, eliminating the need for downstream fine-tuning.
• The model rapidly generates diverse, ready-to-use weights across architectures such as MLPs, ResNets, and Transformers, enhancing ensemble performance and transfer learning.

DeepWeightFlow is a symmetry-aware generative model for neural network weights, applying flow matching directly in weight space to synthesize diverse, high-accuracy neural networks across a broad range of architectures and parameter scales. It addresses the challenges posed by the high dimensionality and symmetries inherent to neural network weights by leveraging canonicalization techniques and a direct vector field regression approach, enabling rapid generation of ready-to-use models without requiring downstream fine-tuning (Gupta et al., 8 Jan 2026).

1. Mathematical Foundations and Weight-Space Flow Matching

DeepWeightFlow is grounded in the flow matching (FM) paradigm. Given a target weight distribution $q(x)$—the empirical distribution of trained network weights—and a source distribution $p_0(x)=\mathcal{N}(0,\sigma^2 I)$, DeepWeightFlow defines linear interpolation paths: $$x_t = (1 - t)x_0 + t x_1, \quad x_0 \sim p_0, \ x_1 \sim q, \ t \in [0,1]$$ The instantaneous path velocity is $u(x_t, t) = x_1 - x_0$. A neural network $v_\theta(x, t)$ is trained via the regression loss: $$\mathcal{L}_{\mathrm{FM}}(\theta) = \mathbb{E}_{t \sim [0, 1], \ (x_0, x_1)} \|v_\theta(x_t, t) - u(x_t, t)\|^2$$ Solving the ODE $\frac{dx}{dt} = v_\theta(x, t), \ x(0) = x_0$ maps the source to the target distribution. The vector field $v_\theta$ thus implements a continuous-time, symmetry-aware transformation in weight space. The associated probability flow satisfies the continuity equation: $$\partial_t p_t(x) + \nabla_x \cdot (p_t(x)v_\theta(x, t)) = 0$$ This approach differs from diffusion- and autoencoder-based generative frameworks by regressing directly in the original (or dimensionally reduced) weight space (Gupta et al., 8 Jan 2026).

2. Canonicalization of Neural Network Weights and Symmetry Handling

Neural network weight spaces are characterized by permutation symmetries (e.g., hidden unit or attention head permutations) and, in certain architectures, scaling symmetries. Such symmetries generate redundant equivalent representations, obstructing efficient generative modeling.

DeepWeightFlow incorporates canonicalization procedures tailored for specific architectures:

• Git Re-Basin [Ainsworth et al. 2023]: Applied to MLPs and ResNets, this procedure solves the Symmetric Orthogonal Bilinear Assignment Problem (SOBLAP) via coordinate descent and the Hungarian algorithm. It finds permutations of adjacent layers to maximize Frobenius inner product alignment with a fixed reference, collapsing permutation modes.
• TransFusion [Rinaldi et al. 2025]: For Transformers (e.g., ViT, BERT), both inter-head and intra-head permutations are resolved. Inter-head alignment matches singular-value spectra, while intra-head matches are performed over the rows/columns of projection matrices.

Canonicalization ensures that each weight configuration corresponds to a single functional representative, allowing the flow-matching model to focus on the intrinsic geometry of the functional manifold rather than confounding symmetry-related redundancies. As confirmed by geometric flow model experiments (Erdogan, 27 Mar 2025), accounting for permutation (and, when possible, scaling) symmetries substantially improves generative efficiency and sample coverage.

3. Model Architecture and Inference Pipeline

The DeepWeightFlow architecture is a fully connected MLP (multi-layer perceptron) conditioned on time:

• The instantaneous time parameter $t$ is embedded via a separate MLP (with LayerNorm, dropout, and GELU), generating $p_0(x)=\mathcal{N}(0,\sigma^2 I)$0.
• The network input is the concatenation of the canonicalized, possibly PCA-reduced, flattened weight vector $p_0(x)=\mathcal{N}(0,\sigma^2 I)$1 and the time embedding.
• The backbone consists of several dense layers (width $p_0(x)=\mathcal{N}(0,\sigma^2 I)$2), each with LayerNorm, GELU, and dropout; the output layer yields $p_0(x)=\mathcal{N}(0,\sigma^2 I)$3.

For high-dimensional cases, incremental or Dual PCA is used to project weights to a lower-dimensional subspace, with the inverse transform applied after sampling. At inference, a sample $p_0(x)=\mathcal{N}(0,\sigma^2 I)$4 is propagated via a fourth-order Runge–Kutta ODE solver ($p_0(x)=\mathcal{N}(0,\sigma^2 I)$5–20), the output is reshaped into layer-wise tensors, and BatchNorm running statistics are recalibrated as needed.

Sampling is thus both rapid and scalable: on modern hardware, even large networks can be generated in seconds per instance, and ensemble generation scales correspondingly. No downstream gradient-based finetuning is necessary for competitive accuracy (Gupta et al., 8 Jan 2026).

4. Experimental Evaluation and Quantitative Performance

DeepWeightFlow was empirically validated across a variety of architectures and datasets, including MLPs, small CNNs, ResNet-18 (CIFAR-10/STL-10), ViT-Small (CIFAR-10), and BERT-Base. In all cases, training utilized 100 independently trained canonicalized models per task, with FM training for 30,000 iterations.

Key results:

• Accuracy without Fine-Tuning: Generated weights closely matched reference solutions.
  • MNIST 3-layer MLP: $p_0(x)=\mathcal{N}(0,\sigma^2 I)$6 vs. $p_0(x)=\mathcal{N}(0,\sigma^2 I)$7.
  • ResNet-18 on CIFAR-10: $p_0(x)=\mathcal{N}(0,\sigma^2 I)$8 (DeepWeightFlow) vs. $p_0(x)=\mathcal{N}(0,\sigma^2 I)$9 (reference, w/ Git Re-Basin).
  • ViT-Small on CIFAR-10: $$x_t = (1 - t)x_0 + t x_1, \quad x_0 \sim p_0, \ x_1 \sim q, \ t \in [0,1]$$0 (DeepWeightFlow, w/ TransFusion) vs. $$x_t = (1 - t)x_0 + t x_1, \quad x_0 \sim p_0, \ x_1 \sim q, \ t \in [0,1]$$1.
• Generation Speed: Substantially outperforms prior methods. For ResNet-18, models are generated in $$x_t = (1 - t)x_0 + t x_1, \quad x_0 \sim p_0, \ x_1 \sim q, \ t \in [0,1]$$21.4s (A100 GPU); ensemble of 100 in $$x_t = (1 - t)x_0 + t x_1, \quad x_0 \sim p_0, \ x_1 \sim q, \ t \in [0,1]$$32 minutes.
• Transfer Learning: Generated networks show strong transfer performance, matching or exceeding FLoWN and random initialization baselines in zero-shot and fine-tuned settings.
• Ensemble Diversity: Ensembles of generated models exhibit high accuracy (statistically comparable to distinct hand-trained models) with considerable diversity (mIoU $$x_t = (1 - t)x_0 + t x_1, \quad x_0 \sim p_0, \ x_1 \sim q, \ t \in [0,1]$$4), facilitating model averaging (Gupta et al., 8 Jan 2026).

5. Comparison with Related Generative Methodologies

DeepWeightFlow distinguishes itself from prior generative models for neural network weights:

• Diffusion Models (e.g., RPG, P-diff, D2NWG): Diffusion-based approaches either generate only partial weights, require hours per sample, or necessitate downstream fine-tuning. DeepWeightFlow’s direct vector field regression eliminates autoregressive or VAE bottlenecks and enables complete weight set generation at orders-of-magnitude higher throughput.
• Latent-Space/Kernel-Density Approaches (e.g., SANE): These require finetuning or sequential layer-wise sampling. DeepWeightFlow enables holistic, symmetry-correct, and fully parallel generation.
• Geometry-Aware Flows: Geometric flow models (Erdogan, 27 Mar 2025) confirm the critical benefit of respecting permutation and scaling symmetries through canonicalization and appropriate network parameterizations (e.g., Relational GNNs, Riemannian manifolds).

A summary of comparative characteristics:

Method	Handles Symmetries	Full Weight Generation	Fine-Tuning Required	Sampling Time (ResNet-18)
DeepWeightFlow	Yes	Yes	No	~1.4s
RPG	Partial	Yes	Yes	O(hours)
P-diff, D2NWG	No/Partial	Partial	Yes	O(minutes-hours)
SANE	No/Partial	Partial	Yes	O(minutes)

6. Limitations and Open Problems

Several challenges and future research avenues are identified (Gupta et al., 8 Jan 2026):

• BatchNorm Handling: Recalibration of BatchNorm statistics is necessary post-generation.
• Conditional/Multiclass Generation: Full-rank conditional weight generation across architectures or semantic classes is limited—PCA-reduced latent spaces allow rudimentary multiclass flows, but with degraded performance.
• Scalability: While plausible up to billion-parameter networks using Dual PCA and high-memory GPUs, this remains to be conclusively demonstrated.
• Open Questions: Trade-offs between explicit canonicalization, symmetry-equivariant generative network architectures, and augmentation schemes in weight space; possible efficiency gains from integrating sparsity or low-rank structure into the generative model.

7. Extensions and Applications

DeepWeightFlow’s framework extends naturally to diverse generative, transfer learning, and diagnostic settings:

• Downstream Initialization: Generated weights consistently accelerate convergence and improve generalization compared to random initializations (Saragih et al., 14 Jul 2025).
• Reward Fine-Tuning: Via adjoint matching, the generative vector field can be adapted to optimize arbitrary downstream metrics, closing significant fractions of the gap to oracle-trained solutions in as few as 150 optimizer steps.
• Covariate-Shift Detection: Ensembles of reward-fine-tuned DeepWeightFlow classifiers outperform Detectron-style baselines in detecting distribution shifts, as measured by TPR and AUROC (Saragih et al., 14 Jul 2025).

This suggests broad utility for meta-modeling, diagnostic, and one-shot learning scenarios, with prospects for meta-foundation models capable of universal, conditional generative weight modeling across families of tasks and architectures.