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

TL;DR

This paper introduces NNiT, a width-agnostic neural network generator that models weights as structured patches, enabling the creation of functional networks across various architectures and generalizing well to unseen topologies.

Contribution

The paper presents a novel width-agnostic neural network generation method using patch-based weight modeling and structural alignment via Graph HyperNetworks.

Findings

Achieves >85% success on unseen architectures in robotics tasks

Enables generation of functional networks across diverse architectures

Outperforms baseline methods in generalization

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 robotics tasks, NNiT achieves >85% success on architecture topologies unseen during training, while baseline approaches fail to generalize.