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

Computer Science > Machine Learning arXiv:2603.00180 (cs) [Submitted on 26 Feb 2026] Title: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 View a PDF of the paper titled NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces, by Jiwoo Kim and 5 other authors View PDF HTML (experimental) 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 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 robotic...