A trained neural network can be represented by a parameter vector in high dimensions. Learning distributions over these vectors enables the generation of new models across various tasks and architectures. A central challenge is permutation symmetry: permuting hidden neurons can produce distant parameter vectors representing the same function. This introduces variations that a generative model must account for when learning from trained networks. Existing methods typically address this using networks derived from a common base model or costly approximate neuron alignment. We instead parameterize a flow-matching velocity field with a permutation-equivariant Graph Meta Network, enabling direct learning from independently trained networks without alignment. Extensive experiments show that our method closely reproduces the joint statistics of accuracy, functional similarity, and weight similarity of independently trained collections, providing evidence of generation beyond checkpoint memorization. A single conditional model also generates task-specific networks on heterogeneous architectures and generalizes to unseen hidden-width configurations. On a tabular domain-shift task, intermediate conditioning produces individual networks with performance comparable to logit ensembles across both domains. Taken together, our results show how permutation equivariance enables learning from diverse collections of independently trained networks without permutation alignment.
Figures & tables
Figure 1: Hidden-unit permutations preserve the network function. Colors track each hidden unit and its incident weights.
Figure 2: Method overview. A shared velocity field vθ(G,wt,t,c) , trained on unaligned networks, generates weights conditioned on architecture G and task c . States wt and yt belong to different architectures. A hidden-unit permutation Pg permutes the trajectory and its velocities, producing functionally equivalent networks (upper pair).
Figure 3: Equivariant velocity network. Top: an ODE solver transforms Gaussian weights into a functional network through repeated velocity evaluations. Bottom left: graph encoding, L message-passing blocks, and parameter-velocity prediction. Bottom right: each block combines time- and task-conditioned AdaLN, message passing, and residual updates.
Figure 4: Generation joint evaluation (conceptual). Memorized networks perform well; nonfunctional networks have low similarity. We seek agreement with training networks in both metrics.
Model
MNIST
CIFAR-10
Acc. (%)
IoU
WCS
JWS ( ↑ )
Acc. (%)
IoU
WCS
JWS ( ↑ )
Training ref.
95.9±0.2
0.559±0.015
0.557±0.010
1.00
79.6±0.4
0.609±0.006
0.609±0.013
1.00
Random init.
10.1±2.4
0.052±0.002
0.138±0.003
0.15
10.1±1.7
0.254±0.012
0.255±0.006
0.31
Perturbation
92.2±2.2
0.403±0.073
0.971±0.001
0.54
74.8±1.7
0.610±0.047
0.995±0.000
0.63
MLP
9.8±3.7
0.052±0.002
0.355±0.009
0.23
10.0±1.6
0.256±0.012
0.146±0.003
0.25
DiT
72.0±6.5
0.138±0.032
0.555±0.010
0.54
43.7±5.5
0.326±0.025
0.553±0.015
0.62
Table 1: Unconditional generation on MNIST and CIFAR-10. Accuracy (%), Max IoU, and WCS report mean ±std across networks; JWS ( ↑ ) averages over reference subsamples. Reference JWS =1 denotes an identity comparison. † indicates aligned training weights. Perturbation reports the highest JWS over tested noise levels ( σ=0.25 for MNIST; 0.10 for CIFAR-10). GMN results include flow-guided refinement. Bold marks the highest JWS among evaluated methods.
Alignment
Acc. (%)
IoU
WCS
JWS ( ↑ )
Unaligned
94.4±0.4
0.469±0.022
0.553±0.008
0.91
Aligned
94.3±0.6
0.464±0.023
0.552±0.008
0.90
Table 2: The equivariant GMN trained on the same MNIST collection, unaligned (numbers from Table 1 ) and aligned via Git Re-Basin. Notation as in Table 1 .
Figure 5: Collection and capacity scaling on mfeat-karhunen . Baselines have 2 M (top) or 10 M (bottom) parameters. Columns report test accuracy, Max error-IoU, Max WCS, and JWS. Dashed lines are collection reference means in the first three columns and JWS =1 in the last. Bands show ±1 standard deviation. † denotes post-hoc alignment via Git Re-Basin.
Figure 6: Heterogeneous generation and zero-shot width transfer. A single velocity network with shared parameters generates networks across all 20 tasks. Base architectures list layer widths from input to output; ∗ marks architectures shared by two tasks. Narrow and Wide multiply all hidden widths by 0.5× and 1.5× , respectively, while preserving input and output dimensions. Reference reports the mean accuracy of independently trained networks at the base architecture. Generated accuracies are mean ±std (%) over 100 networks per entry. The final plot averages the task means.
Figure 7: ACS domain conditioning. Faint points show training collections. (a) Endpoint generation. (b) Intermediate conditions c∈{0.4,0.5,0.6} and AUC comparison curves. Panels (a,b) use 50 networks per condition, with guidance scales 1.5 and 4.0 , respectively. Comparison curves average 50 CA–PR checkpoint pairs; Re-Basin labels denote matching iterations. (c) JWS against both domain references, using 100 generated networks at each c∈{0,0.1,…,1} .
Appendix figures & tables15 assets
Supplementary material from the paper’s appendix.
Appendix
(a) Held-out network sanity check
Dataset
MNIST
CIFAR-10
mfeat-karhunen
M
100
100
100
500
1,000
JWS ↑
0.990
0.995
0.981
0.989
0.991
Appendix
Table 3: JWS validation. (a) Held-out networks use the original architectures and training recipes with disjoint seeds, evaluated using the original reference clouds and scales. M is the number of held-out networks evaluated. (b) Full uses both normalizations; Raw removes coordinate scaling; No 3 removes only the final distance normalizer. All variants retain clipping at zero. † denotes aligned training weights. Bold marks the highest score among evaluated methods.
Setting
MNIST MLP
CIFAR-10 CNN
Collection size
10K
50K
Architecture
MLP 784→32→32→10 , ReLU
Three 3×3 convolutions, channels (16,32,48) , padding 1 , ReLU
Pooling and output
Linear classifier
2×2 max-pooling after the first two convolutions; global average pooling; linear classifier
Parameters
26,506
19,450
Optimizer
Adam
SGD; Nesterov momentum 0.9 ; weight decay 5×10−4
Learning rate
10−3 , constant
0.1 , cosine decay to zero
Appendix
Table 4: Classifier architectures and training settings. All affine layers include biases.
Figure 8: Input-layer folding. First-layer weights and biases are encoded as first-hidden-layer node features. Hidden-unit permutations reorder these features while leaving input coordinates fixed; shared encoders and decoders therefore preserve permutation equivariance.
Figure 9: CNN kernel edge representation. One edge carries a vectorized k×k kernel.
Setting
MNIST
CIFAR-10
Architecture
Message-passing blocks
8
16
Node / edge / time dimensions
128 / 128 / 128
192 / 384 / 128
Positional encodings
Orbit-shared
Orbit-shared
Weight normalization
Parameter-group
Parameter-group
Message aggregation
Sum
Sum
Appendix
Table 5: GMN configurations for unconditional generation. Refinement uses corruption time t∗ and k cycles.
Sampling
Acc. (%)
Max IoU
WCS
JWS
MNIST (t∗,k)=(0.75,10)
Direct
92.58±0.89
0.383±0.029
0.433±0.007
0.776
Refined
94.40±0.41
0.469±0.022
0.553±0.008
0.906
CIFAR-10 (t∗,k)=(0.3,10)
Direct
72.11±1.54
0.496±0.019
0.601±0.016
0.879
Refined
75.34±0.80
0.533±0.012
0.630±0.011
0.918
Appendix
Table 6: GMN refinement ablation. Each setting contains 100 generated networks; MNIST uses the same networks before and after refinement. Accuracy, Max IoU, and WCS report mean ±std across networks; JWS uses repeated reference subsampling.
Dataset
Max IoU
WCS
MNIST
0.418±0.016
0.529±0.010
CIFAR-10
0.506±0.009
0.610±0.014
Appendix
Table 7: Leave-one-out maximum similarity among generated networks, reported as mean ±std .
Baseline
MNIST
CIFAR-10
Search coverage
MLP
1
1
Dataset-specific fixed architectures with 21.4 M and 80.4 M parameters, respectively.
DiT
5
2
MNIST: widths {256,368} , learning rates {10−3,5×10−4,3×10−4} , constant/cosine schedules, and EMA decays {0.999,0.9999} in five tested combinations. CIFAR-10: two width- 864 configurations.
DWF
72
9
MNIST: full grid of 3 batch sizes, 3 learning rates, 4 widths, and 2 alignment settings. CIFAR-10: 9 combinations of 3 learning rates, 2 batch sizes, and 2 alignment settings.
SANE
14
4
Seven MNIST widths ( 2.0 – 202.7 M parameters) and two CIFAR-10 widths ( 79.7 – 202.7 M), each evaluated with KDE and Gaussian priors.
P-diff
1
1
Dataset-specific VAE/DDPM configurations, with 23.1 M and 42.3 M total parameters, respectively.
Total
93
17
110 evaluated configurations.
Appendix
Table 8: Baseline search coverage. Counts denote evaluated configurations, including generation-prior choices. The MNIST DWF search is a full Cartesian grid; the exact DiT and CIFAR-10 DWF combinations appear in Table .
Control
Acc. (%)
Max IoU
WCS
JWS
MNIST
Random init.
10.11±2.39
0.052±0.002
0.138±0.003
0.1457±0.0010
σ=0.05
95.81±0.25
0.855±0.029
0.999±0.000
0.4482±0.0011
σ=0.10
95.52±0.35
0.717±0.047
0.995±0.000
0.5155±0.0010
σ=0.25
92.17±2.15
0.403±0.073
0.971±0.001
0.5366±0.0010
σ=0.50
76.24±7.42
0.166±0.044
0.896±0.003
0.4471±0.0012
Appendix
Table 9: Random initialization and perturbation controls. Entries report mean ± std across networks, or reference subsamples for JWS. Bold marks the highest JWS per dataset.
Configuration
Params (M)
Acc. (%)
Max IoU
WCS
JWS
MLP: MNIST
h=352,L=8dt=128
21.4
9.78±3.74
0.0517±0.0020
0.3554±0.0086
0.2335±0.0010
MLP: CIFAR-10
h=1184,L=10dt=256
80.4
9.98±1.55
0.2558±0.0120
0.1464±0.0025
0.2520±0.0009
DiT: MNIST
h=256,η=10−3 Cosine; ρ=0.999
8.2
9.95±5.40
0.0517±0.026
0.5345±0.009
0.2625±0.0010
Appendix
Table 10: Consolidated baseline configurations and sweep results. h denotes hidden or transformer width, η learning rate, ρ EMA decay. b batch size, L the number of blocks, and dt the time-embedding dimension. Parameter counts are approximate and given in millions. Entries use x±s≡x±s ; accuracy, Max IoU, and WCS summarize generated networks. JWS is reported as a mean with subsampling standard deviation. Bold marks selected configurations reported in Table 1 . The 36 unaligned MNIST DWF runs are summarized by per-metric extrema, which need not correspond to the same configuration.
Table 11: Hyperparameters for the scaling experiment. Width denotes the GMN feature dimension, DiT hidden dimension, or DWF hidden width. Orbit-shared encodings use hidden-unit permutation orbits.
N
Model
Acc. (%)
Max IoU
WCS
JWS
100
GMN
62.25±11.81
0.519±0.067
0.792±0.051
0.721
DiT
89.24±2.47
0.979±0.025
1.000±0.000
0.564
DiT-L
86.24±13.33
0.974±0.101
0.988±0.069
0.558
DWF
85.15±17.45
0.860±0.117
0.971±0.120
0.605
DWF †
89.39±2.33
0.906±0.068
0.998±0.006
0.600
DWF-L
88.69±3.03
0.897±0.076
0.998±0.002
0.603
Appendix
Table 12: Scaling results for all seven configurations. Accuracy, Max IoU, and WCS report x±s≡x±s across 100 generated networks; JWS measures joint distributional agreement with the fixed reference cloud. Models have approximately 2 M parameters, or 10 M for the -L variants. † denotes Git Re-Basin alignment of the training collection. Bold marks the highest JWS at each N .
Data and training collections
Data source
Folktables ACSIncome ( Ding et al., 2021 ) ; 2018 one-year ACS data
Domains and split
CA ( c=0 ), PR ( c=1 ); up to 30K examples per domain; stratified 80/20 training/test split
Categorical features
One-hot encode COW , MAR , RELP , SEX , and RAC1P
Numerical features
Median-impute and standardize AGEP , SCHL , and WKHP
Preprocessing
Fitted jointly on both training sets, then frozen; 44 input features
Collection size
5K independently trained networks per domain; default PyTorch initialization; final checkpoints only
Appendix
Table 13: ACS data, training collections, conditional flow, and AUC evaluation settings.
Batch size / epochs: 128/25 ; 32/300 for balance-scale, hayes-roth, and solar-flare (OpenML 40686)
Conditional flow
Message-passing blocks
10
Parameters (M)
15.13
Appendix
Table 14: Heterogeneous collection construction, shared flow training, and evaluation. All datasets and width configurations use the same flow checkpoint and sampling settings. All the networks in the training collections were initialized with different seeds.
Width
Scratch
Generated
Δ
kr-vs-kp 73→32→32→2
0.5×
97.38±0.72
96.02±1.62
−1.36
1×
98.61±0.51
97.62±0.60
−0.99
1.5×
99.13±0.39
97.13±0.94
−2.00
letter 16→32→32→26
0.5×
79.54±0.80
54.38±3.51
−25.16
Appendix
Table 15: Heterogeneous generation and width transfer. Headings give each dataset’s native architecture; width multipliers apply to all hidden layers, with 0.5× and 1.5× unseen during flow training. Accuracy (%) is mean ±std ; Δ is generated minus scratch accuracy in percentage points. All generated networks use one shared flow. Scratch references contain 5K native networks or 100 networks per variant; each generated entry uses 100 networks. All evaluations use the same held-out split per dataset. The final block averages all 20 datasets equally. Solar-flare (5 classes) is OpenML 40686 .
In this work, we develop theoretical foundation for flow matching with neural-network-parameterized conditional velocity fields. We establish convergence guarantees for gradient descent in the over-parameterized 2-layered ReLU neural network regime. We derive generalization bounds for the conditional velocity-field matching objective. Building on these results, we provide Wasserstein-distance guarantees for the samples generated by the induced flow. Our analysis is based on generalization bound for multi-task representation learning with unbounded losses, which may be of independent interest beyond flow-based generative modeling. These theoretical results are validated through extensive experiments on both synthetic and real-world image benchmarks.
Yihan He, Qishuo Yin, Yuan Cao +2
Princeton University, Princeton, NJ, USA. · The University of Hong Kong, Hong Kong. · Northwestern University, Evanston, IL, USA.
Generative modeling of neural network parameters is often tied to architectures because standard parameter representations rely on known weight-matrix dimensions. Generation is further complicated by permutation symmetries that allow networks to model similar input-output functions while having widely different, unaligned parameterizations. In this work, we introduce Neural Network Diffusion Transformers (NNiTs), which generate weights in a width-agnostic manner by tokenizing weight matrices into patches and modeling them as locally structured fields. We establish that Graph HyperNetworks (GHNs) with a convolutional neural network (CNN) decoder structurally align the weight space, creating the local correlation necessary for patch-based processing. Focusing on Multilayer Perceptrons (MLPs), where permutation symmetry is especially apparent, NNiTs generate fully functional networks across a range of architectures. Our approach jointly models discrete architecture tokens and continuous weight patches within a single sequence model. On ManiSkill3 robotics tasks, NNiT achieves >85% success on architecture topologies unseen during training, while baseline approaches fail to generalize; the same pipeline also generalizes to MNIST classification beyond the robotic control setting.
Jiwoo Kim, Swarajh Mehta, Hao-Lun Hsu +3
Department of Electrical and Computer Engineering, Duke University, NC, USA · Department of Computer Science, Duke University, NC, USA · Computer Science & Artificial Intelligence Laboratory, Massachusetts Institute of Technology, MA, USA
Learning permutations is fundamental to sorting, ranking, and matching, but existing differentiable methods based on entropy-regularized Sinkhorn produce a single softened solution and collapse under ambiguity. We present PermFlow, a conditional flow matching framework that operates directly on the affine subspace of matrices with unit row and column sums. A closed-form tangent-space projector preserves these constraints exactly along every trajectory, by construction rather than through iterative correction, and a nearest-target coupling routes distinct noisy initializations toward distinct valid permutations. The result is a model that captures multimodal permutation distributions rather than collapsing them to a single mode. On a visual sorting task with blended-digit ambiguity and a symmetric linear assignment problem, PermFlow achieves high accuracy on unambiguous inputs and recovers both valid permutations under ambiguity, where Sinkhorn-based baselines structurally fail.
Yimeng Min, Carla P. Gomes
Department of Computer Science Cornell University Ithaca, NY, USA