NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces
Authors: Jiwoo Kim, Swarajh Mehta, Hao-Lun Hsu, Hyunwoo Ryu, Yudong Liu, Miroslav Pajic
Organizations: 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
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.
Figures & tables
Figure 1: Width-Agnostic Synthesis via Multimodal Tokenization . Unlike existing models, the NNiT decouples functional logic from fixed matrix dimensions, allowing the zero-shot synthesis of optimal weights for architectural topologies entirely unseen during training.
Figure 2: NNiT Framework Overview. Left: Unified Generative Architecture. We formulate neural synthesis as a multimodal sequence task. Discrete architecture tokens (orange) and continuous weight matrices (blue) are unified into a single sequence, with weights processed as spatially correlated patches. A Diffusion Transformer (DiT) models the joint distribution using per-modality timestep conditioning ( μa,Σa,μw,Σw ) via the Mixture of Noise Levels (MoNL) framework, enabling both co-design p(a,w) and conditional synthesis p(w∣a) . Right: Deployment Pipeline. During inference, sampled architecture tokens are decoded into layer widths. The generated weight tensors are then extracted to match these target dimensions, assembling a directly executable MLP.
Figure 3: Visualizing Structural Alignment and Induced Geometry. Comparison of the weight magnitude profiles across 35 independent seeds. Top (GHN): The consistent alignment across the seeds demonstrates that the GHN 1) successfully spatially aligns the weight spaces , effectively resolving the permutation ambiguity inherent in neural networks. Moreover, the visible structural banding indicates that the GHN 2) imposes meaningful geometric structure (spatial correlation), transforming independent parameters into a spatially aligned space. Bottom (SGD): In contrast, SGD weights exhibit unstructured noise due to arbitrary permutations, lacking both structural alignment and local spatial geometry.
Metric
GHN
SGD
Policy Performance
Mean Return
39.35
39.49
Mean Success Rate
99.0%
98.9%
Weight Diversity (Nearest Neighbor)
L2 Distance ( ↑ )
54.66
41.73
Cosine Similarity ( ↓ )
0.07
0.02
Table 1: Alignment vs. Performance. Evaluation of 35 independent seeds. Both methods achieve optimal performance, yet GHNs induce the structural alignment necessary for tokenization without sacrificing weight diversity, confirming the absence of mode collapse.
PickCube-v1
PushCube-v1
StackCubeEasy-v1
Architecture
Return
Success
Return
Success
Return
Success
Dataset
39.37±0.55
98%
42.02±0.4
100%
39.53±0.78
94%
SANE
Seen
3.81±0.29
1%
6.90±0.25
6%
4.64±0.16
0%
D2NWG
38.59±1.35
98%
42.28±0.05
100%
37.65±0.69
90%
NNiT (Ours)
38.77±0.41
98%
42.10±0.16
100%
38.38±0.51
91%
SANE
3.65±0.41
1%
5.89±0.32
2%
3.71±0.13
0%
Table 2: Architecture-conditioned weight generation ( p(w∣a) ). We compare NNiT against baseline methods on both seen and unseen (zero-shot) architectural configurations. Reporting the top-10 performance metric, we observe that while the gap is negligible on training configurations, performance diverges sharply on unseen topologies. Baselines degrade significantly, while NNiT maintains robust performance, validating its effectiveness as a width-agnostic neural network generator.
PickCube-v1
PushCube-v1
StackCubeEasy-v1
Return
Success
Return
Success
Return
Success
NNiT(Ours)
38.85±0.32
99%
41.99±0.10
100%
37.57±0.85
90%
Table 3: Multimodal Joint Synthesis ( p(a,w) ). We evaluate NNiT’s capacity to spontaneously synthesize complete functional policies. As no baselines support joint synthesis, NNiT is presented in isolation. The results demonstrate near-perfect success rates ( 99%–100% ) across PickCube and PushCube and 90% for StackCubeEasy. This confirms that the model can generate high-performing policies without a fixed architectural prompt.
Policy Architecture
Partition
Return
Success
Training (Seen) Topologies
32, 16, 64, 32
Seen
39.69
100%
32, 64, 16, 32
Seen
39.01
100%
32, 32, 32, 16
Seen
38.99
100%
64, 64, 16, 32
Seen
38.81
98%
Zero-Shot (Unseen) Topologies
Table 4: Architecture Diversity & Generalization. Specific functional networks sampled from the generated distribution. Note that policy architectures list hidden layer widths only, omitting input and output layers. The bottom section highlights Zero-Shot Synthesis on topologies entirely unseen (Test Set) during training. Notably, NNiT achieves 98% success on the unseen configuration [input,32,16,16,16,output] .
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Environment
Mean NN L2 ( ↑ )
Mean NN Cosine ( ↓ )
PickCube-v1
46.94
0.17
PushCube-v1
49.66
0.12
StackCubeEasy-v1
72.79
0.21
Appendix
Table 5: Quantitative Diversity Metrics. Evaluation of 100 filtered policy weights. The high L2 distances and low cosine similarities confirm that the dataset maintains significant parametric diversity despite the structural alignment.
Figure 4: Dataset Diversity Analysis. Histograms of pairwise L2 distances and Cosine Similarities across all three environments. The consistently high L2 distances and low cosine similarities confirm that the structural alignment imposed by the GHN does not result in mode collapse.
Figure 5: Visualization of Topological Anchoring. Heatmaps of neuron-wise weight magnitudes for 3 selected architectures across 100 filtered seeds. The vertical banding visually confirms the induction of spatial correlation, validating the premise that these weights can be treated as continuous fields.
CKA
GHN
SGD
Overall
0.849±0.088
0.933±0.030
Layer 1
0.917±0.016
0.957±0.005
Layer 2
0.880±0.033
0.926±0.022
Layer 3
0.751±0.083
0.918±0.037
Appendix
Table 6: Representational Diversity via CKA. Lower CKA indicates greater functional diversity between independently trained networks. GHN-generated networks exhibit lower CKA than SGD-trained networks at every layer, with the gap widening in deeper layers.
Environment Specifics
Parameter
PickCube-v1
PushCube-v1
StackCubeEasy-v1
GHN Training
Iterations
500
500
1500
Beta Decay
0.97
0.97
0.98
Buffer Size
1000000
1000000
1000000
Meta Batch Size
8
8
16
Appendix
Table 7: Hyperparameters for Dataset Generation and Filtering
Hyperparameter
PickCube
PushCube
StackCubeEasy
Diffusion Process
Timesteps ( T )
1000
1000
1000
Noise Schedule
Linear
Linear
Linear
Prediction Target
Noise ( ϵ )
Noise ( ϵ )
Noise ( ϵ )
Optimization
Optimizer
AdamW
AdamW
AdamW
Appendix
Table 8: Hyperparameter Configuration. Training and optimization settings for the diffusion process.
Metric
Method
PickCube-v1
PushCube-v1
StackCubeEasy-v1
Mean (80 policies)
NNiT (Ours)
78.8%
74.7%
54.9%
D2NWG
15.6%
60.0%
12.2%
Top-10
NNiT (Ours)
98.8%
100%
86.4%
D2NWG
59.4%
89.2%
41.8%
Worst Architecture
NNiT (Ours)
64.8%
59.6%
42.6%
D2NWG
5.6%
41.2%
2.0%
Appendix
Table 9: Full Distribution on Unseen Architectures. Mean (over all 80 generated policies), Top-10, and Worst-Arch. NNiT’s worst single architecture outperforms D2NWG’s overall mean on both PickCube-v1 and StackCubeEasy-v1 .
Figure 6: Parametric Diversity Analysis. Pairwise Euclidean distance (left) and cosine similarity (right) matrices for the top 10 generated policies. The dominance of high distances and low cosine similarities confirms that the generative model preserves variance within the parameter space, avoiding mode collapse.
Figure 7: PickCube-v1 Policy Rollout. A sequential visualization of policies generated by NNiT across diverse topologies. Left (Y-axis): Target architecture configurations denoted by their hidden layer widths (e.g., 32-16-64-32). Top (X-axis): Temporal snapshots of the rollout in 10-step increments. The agent successfully navigates to the target object, executes a stable grasp, and lifts the cube to the designated goal, visually demonstrating that the sampled NNiT MLPs operate correctly.
Figure 8: PushCube-v1 Policy Rollout. A sequential visualization of policies generated by NNiT across diverse topologies. Left (Y-axis): Target architecture configurations denoted by their hidden layer widths (e.g., 32-64-32-64). Top (X-axis): Temporal snapshots of the rollout in 10-step increments. The policy effectively coordinates the end-effector to manipulate the object, pushing it into the target goal region.
Figure 9: StackCubeEasy-v1 Policy Rollout. A sequential visualization of policies generated by NNiT across diverse topologies. Left (Y-axis): Target architecture configurations denoted by their hidden layer widths (e.g., 64-64-64-32). Top (X-axis): Temporal snapshots of the rollout in 10-step increments. The synthesized agent demonstrates high-precision control by successfully grasping the red cube and placing it on top of the green goal cube.
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.
Arkadi Piven, Yam Eitan, Guy Bar-Shalom +5
Technion – Israel Institute of Technology · Technical University of Munich · Munich Center for Machine Learning +1
Recent advances in robot learning have driven rapid progress in humanoid control, enabling increasingly agile behaviors in the real world. However, humanoid control networks are typically obtained through conventional training procedures, while comparatively little attention has been given to directly modeling the distribution of their learned weights and generating functional networks in weight space. In this work, we investigate whether high-performing humanoid control networks can be generatively modeled directly through their weights. We adopt a latent diffusion framework that couples a linear autoencoder with a Diffusion Transformer. The autoencoder compresses actor network parameters into a compact latent space, and the Diffusion Transformer models the distribution of these latents. Trained on networks collected from reinforcement learning for the Unitree G1 robot, our approach generates deployable networks that perform various motions in the real world, including walking, dancing, and spin kick. Beyond generation, we further study the learned weight space to better understand the structure captured by these humanoid control networks. We identify a structural principle in the first-layer weights that consistently reflects the robot's kinematic prior. Notably, this kinematic prior emerges under both Reinforcement Learning (RL) and Imitation Learning (IL) training, suggesting that it is not specific to a particular optimization process but may instead reflect common knowledge associated with humanoid embodiment. Together, these results demonstrate the feasibility of generative modeling in humanoid network weight space, while providing initial insights into the information encoded by their learned weights. Additional results are available on our project website: https://humanoid-weight-diffusion.github.io/
Huixuan Cheng, Bike Zhang, Shankar Sastry +1
The authors are with the University of California, Berkeley.
Neural network weights are typically viewed as the end product of training, while most deep learning research focuses on data, features, and architectures. However, recent advances show that the set of all possible weight values (weight space) itself contains rich structure: pretrained models form organized distributions, exhibit symmetries, and can be embedded, compared, or even generated. Understanding such structures has tremendous impact on how neural networks are analyzed and compared, and on how knowledge is transferred across models, beyond individual training instances. This emerging research direction, which we refer to as Weight Space Learning (WSL), treats neural weights as a meaningful domain for analysis and modeling. This survey provides the first unified taxonomy of WSL. We categorize existing methods into three core dimensions: Weight Space Understanding (WSU), which studies the geometry and symmetries of weights; Weight Space Representation (WSR), which learns embeddings over model weights; and Weight Space Generation (WSG), which synthesizes new weights through hypernetworks or generative models. We further show how these developments enable practical applications, including model retrieval, continual and federated learning, neural architecture search, and data-free reconstruction. By consolidating fragmented progress under a coherent framework, this survey highlights weight space as a learnable, structured domain with growing impact across model analysis, transferring, and weight generation. We release an accompanying resource at https://github.com/Zehong-Wang/Awesome-Weight-Space-Learning.
Xiaolong Han, Zehong Wang, Bo Zhao +8
University of Surrey · University of Notre Dame · University of California San Diego +4