Reproduce “NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces” on OpenResearch

About this paper

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

Researchers at Duke University and MIT developed NNiT, a multimodal diffusion transformer capable of generating neural network parameters independent of layer width by representing weights as patch-based continuous fields. The framework leverages Graph HyperNetworks to induce structural alignment in weight spaces, enabling zero-shot synthesis of functional MLP policies for unseen architectures with high success rates in robotic control tasks and MNIST classification.

Key insights

• Graph HyperNetworks (GHNs) with a CNN decoder can inherently generate structurally aligned weight spaces, effectively collapsing permutation symmetries into a consistent topological structure.
• Treating neural network weights as a continuous spatial field allows for patch-based tokenization, facilitating width-agnostic generation and zero-shot synthesis for novel architectural topologies.
• A multimodal Diffusion Transformer can effectively capture interdependencies between discrete architecture tokens and continuous weight patches, enabling joint and conditional synthesis of functional networks.

Summary by alphaXiv.