NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces
This work addresses the challenge of generating neural network parameters independent of architecture width and permutation symmetries, which is a significant problem for researchers and practitioners in generative modeling of neural networks.
This paper introduces Neural Network Diffusion Transformers (NNiTs), a method for generating neural network weights in a width-agnostic manner by tokenizing weight matrices into patches. NNiT successfully generates functional MLPs and achieves over 85% success on unseen architecture topologies in ManiSkill3 robotics tasks, where baseline methods fail.
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 MLPs, where permutation symmetry is especially apparent, NNiT generates 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.