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

arXiv: https://arxiv.org/abs/2603.00180 · HTML: https://mlsanity.com/paper/2603.00180/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 55 references that are in the map.
- First posted on arXiv 26 Feb 2026.

## Builds on (most cited first, 30 of 44)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Structured 3D Latents for Scalable and Versatile 3D Generation](https://mlsanity.com/paper/2412.01506/) (2024)
- [Latte: Latent Diffusion Transformer for Video Generation](https://mlsanity.com/paper/2401.03048/) (2024)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model](https://mlsanity.com/paper/2408.11039/) (2024)
- [MDTv2: Masked Diffusion Transformer is a Strong Image Synthesizer](https://mlsanity.com/paper/2303.14389/) (2023)
- [ManiSkill3: GPU Parallelized Robotics Simulation and Rendering for Generalizable Embodied AI](https://mlsanity.com/paper/2410.00425/) (2024)
- [MM-Diffusion: Learning Multi-Modal Diffusion Models for Joint Audio and Video Generation](https://mlsanity.com/paper/2212.09478/) (2022)
- [Direct3D: Scalable Image-to-3D Generation via 3D Latent Diffusion Transformer](https://mlsanity.com/paper/2405.14832/) (2024)
- [SANA 1.5: Efficient Scaling of Training-Time and Inference-Time Compute in Linear Diffusion Transformer](https://mlsanity.com/paper/2501.18427/) (2025)
- [HyperDiffusion: Generating Implicit Neural Fields with Weight-Space Diffusion](https://mlsanity.com/paper/2303.17015/) (2023)
- [Lumina-T2X: Transforming Text into Any Modality, Resolution, and Duration via Flow-based Large Diffusion Transformers](https://mlsanity.com/paper/2405.05945/) (2024)
- [Graph HyperNetworks for Neural Architecture Search](https://mlsanity.com/paper/1810.05749/) (2018)
- [Protein structure generation via folding diffusion](https://mlsanity.com/paper/2209.15611/) (2022)
- [DiT-3D: Exploring Plain Diffusion Transformers for 3D Shape Generation](https://mlsanity.com/paper/2307.01831/) (2023)
- [Equivariant Architectures for Learning in Deep Weight Spaces](https://mlsanity.com/paper/2301.12780/) (2023)
- [Snap Video: Scaled Spatiotemporal Transformers for Text-to-Video Synthesis](https://mlsanity.com/paper/2402.14797/) (2024)
- [SANA-Video: Efficient Video Generation with Block Linear Diffusion Transformer](https://mlsanity.com/paper/2509.24695/) (2025)
- [Parameter Prediction for Unseen Deep Architectures](https://mlsanity.com/paper/2110.13100/) (2021)
- [JavisDiT: Joint Audio-Video Diffusion Transformer with Hierarchical Spatio-Temporal Prior Synchronization](https://mlsanity.com/paper/2503.23377/) (2025)
- [Learning to Learn with Generative Models of Neural Network Checkpoints](https://mlsanity.com/paper/2209.12892/) (2022)
- [FiT: Flexible Vision Transformer for Diffusion Model](https://mlsanity.com/paper/2402.12376/) (2024)
- [PixelDiT: Pixel Diffusion Transformers for Image Generation](https://mlsanity.com/paper/2511.20645/) (2025)
- [All-atom Diffusion Transformers: Unified generative modelling of molecules and materials](https://mlsanity.com/paper/2503.03965/) (2025)
- [Diffusion-EDFs: Bi-equivariant Denoising Generative Modeling on SE(3) for Visual Robotic Manipulation](https://mlsanity.com/paper/2309.02685/) (2023)
- [Graph Diffusion Transformers for Multi-Conditional Molecular Generation](https://mlsanity.com/paper/2401.13858/) (2024)

## Built on by (most cited first, 3 of 3)

- [Universal Hypernetworks for Arbitrary Models](https://mlsanity.com/paper/2604.02215/) (2026)
- [Permutation-Equivariant Flow Matching for Alignment-Free Neural Weight Generation](https://mlsanity.com/paper/2609.32833/) (2026)
- [SLEM: Spectrum-Aware Latent Expert Merging](https://mlsanity.com/paper/2610.16484/) (2026)

## Abstract (arXiv)

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.