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

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.

Modern AI systems run on neural networks whose behavior is set by numerous weights, which are normally fixed through slow, expensive training. An alternative is to build a second AI that generates these weights directly. The obstacle is that existing weight-generators are locked to a fixed layer width: they can often be made deeper, but changing how wide the layers are makes them fail. This happens because the same network can be written in countless reshuffled ways, so the weights look like structureless noise, and because prior methods flatten each network into a fixed-length list of numbers whose length is tied to one specific width. Our key finding is that Graph HyperNetworks (GHNs) — models that generate other networks' weights — naturally lay those weights out in a consistent structure instead of the usual scrambled mess. Because the weights are already aligned, our model, NNiT, can generate them directly. And because NNiT's generation process is not tied to any fixed network width, it generalizes to unseen widths, producing fully working robot-control and digit-recognition networks where previous approaches collapse.