NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces

Jiwoo Kim; Swarajh Mehta; Hao-Lun Hsu; Hyunwoo Ryu; Yudong Liu; Miroslav Pajic

cs.LGFeb 26, 2026

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

Abstract

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 ( $\mu_{a},\Sigma_{a},\mu_{w},\Sigma_{w}$ ) via the Mixture of Noise Levels (MoNL) framework, enabling both co-design $p(\mathbf{a},\mathbf{w})$ and conditional synthesis $p(\mathbf{w}|\mathbf{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.

Policy Performance
Metric	GHN	SGD
Mean Return	$39.35$	$39.49$
Mean Success Rate	$99.0\%$	$98.9\%$
Weight Diversity (Nearest Neighbor)
$L_{2}$ Distance ( $\uparrow$ )	$54.66$	$41.73$
Cosine Similarity ( $\downarrow$ )	$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\pm 0.55$	$98\%$	$42.02\pm 0.4$	$100\%$	$39.53\pm 0.78$	$94\%$
SANE	Seen	$3.81\pm 0.29$	$1\%$	$6.90\pm 0.25$	$6\%$	$4.64\pm 0.16$	$0\%$
D2NWG		$38.59\pm 1.35$	$98\%$	$42.28\pm 0.05$	$100\%$	$37.65\pm 0.69$	$90\%$
NNiT (Ours)		$38.77\pm 0.41$	$98\%$	$42.10\pm 0.16$	$100\%$	$38.38\pm 0.51$	$91\%$
SANE		$3.65\pm 0.41$	$1\%$	$5.89\pm 0.32$	$2\%$	$3.71\pm 0.13$	$0\%$

Table 2: Architecture-conditioned weight generation (

p(\mathbf{w}|\mathbf{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\pm 0.32$	$99\%$	$41.99\pm 0.10$	$100\%$	$37.57\pm 0.85$	$90\%$

Table 3: Multimodal Joint Synthesis (

p(\mathbf{a},\mathbf{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\%\text{--}100\%

) across PickCube and PushCube and

90\%

for StackCubeEasy. This confirms that the model can generate high-performing policies without a fixed architectural prompt.

Training (Seen) Topologies
Policy Architecture	Partition	Return	Success
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

[\text{input},32,16,16,16,\text{output}]

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Environment	Mean NN $L_{2}$ ( $\uparrow$ )	Mean NN Cosine ( $\downarrow$ )
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

L_{2}

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 $L_{2}$ distances and Cosine Similarities across all three environments. The consistently high $L_{2}$ 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\pm 0.088$	$0.933\pm 0.030$
Layer 1	$0.917\pm 0.016$	$0.957\pm 0.005$
Layer 2	$0.880\pm 0.033$	$0.926\pm 0.022$
Layer 3	$0.751\pm 0.083$	$0.918\pm 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

Diffusion Process
Hyperparameter	PickCube	PushCube	StackCubeEasy
Timesteps ( $T$ )	1000	1000	1000
Noise Schedule	Linear	Linear	Linear
Prediction Target	Noise ( $\epsilon$ )	Noise ( $\epsilon$ )	Noise ( $\epsilon$ )
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)	$\mathbf{78.8\%}$	$\mathbf{74.7\%}$	$\mathbf{54.9\%}$
Mean (80 policies)	D2NWG	$15.6\%$	$60.0\%$	$12.2\%$
Top-10	NNiT (Ours)	$\mathbf{98.8\%}$	$\mathbf{100\%}$	$\mathbf{86.4\%}$
Top-10	D2NWG	$59.4\%$	$89.2\%$	$41.8\%$
Worst Architecture	NNiT (Ours)	$\mathbf{64.8\%}$	$\mathbf{59.6\%}$	$\mathbf{42.6\%}$
Worst Architecture	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.

Explore similar work

Permutation-Equivariant Flow Matching for Alignment-Free Neural Weight Generation
Sep 26, 2026Arkadi Piven, Yam Eitan, Guy Bar-Shalom +5Flow Matching Equivariant Neural Networks
Technion – Israel Institute of Technology · Technical University of Munich · Munich Center for Machine Learning +1
Modeling and Understanding Neural Weights of Humanoid Policies with Diffusion
Oct 14, 2026Huixuan Cheng, Bike Zhang, Shankar Sastry +1Generative Modeling Humanoid Robot Control
The authors are with the University of California, Berkeley.
A Survey of Weight Space Learning: Understanding, Representation, and Generation
Mar 10, 2026Xiaolong Han, Zehong Wang, Bo Zhao +8Weight-Space Analysis
University of Surrey · University of Notre Dame · University of California San Diego +4

Sep 26, 2026J/K move · Enter open · S save

Permutation-Equivariant Flow Matching for Alignment-Free Neural Weight Generation

Arkadi Piven, Yam Eitan, Guy Bar-Shalom, Fabrizio Frasca+4

Technion – Israel Institute of Technology · Technical University of Munich · Munich Center for Machine Learning+1

Flow Matching Equivariant Neural Networks

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Permutation-Equivariant Flow Matching for Alignment-Free Neural Weight Generation

Modeling and Understanding Neural Weights of Humanoid Policies with Diffusion

A Survey of Weight Space Learning: Understanding, Representation, and Generation