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

• URL: http://arxiv.org/abs/2603.00180v1
• Date: Thu, 26 Feb 2026 20:47:30 GMT
• 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,
• Abstract summary: We introduce Neural Network Diffusion Transformers (NNiTs), which generate weights in a width-agnostic manner.<n>On ManiSkill3 robotics tasks, NNiT achieves >85% success on architecture topologies unseen during training, while baseline approaches fail to generalize.
• Score: 15.631276865948097
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• SWAT-NN: Simultaneous Weights and Architecture Training for Neural Networks in a Latent Space [6.2241272327831485]
  We propose a framework that simultaneously optimize both the architecture and the weights of a neural network.<n>Our framework first trains a universal multi-scale autoencoder that embeds both architectural and parametric information into a continuous latent space.<n>Given a dataset, we then randomly initialize a point in the embedding space and update it via gradient descent to obtain the optimal neural network.
  arXiv  Detail & Related papers  (2025-06-09T22:22:37Z)
• Enhancing lattice kinetic schemes for fluid dynamics with Lattice-Equivariant Neural Networks [79.16635054977068]
  We present a new class of equivariant neural networks, dubbed Lattice-Equivariant Neural Networks (LENNs) Our approach develops within a recently introduced framework aimed at learning neural network-based surrogate models Lattice Boltzmann collision operators. Our work opens towards practical utilization of machine learning-augmented Lattice Boltzmann CFD in real-world simulations.
  arXiv  Detail & Related papers  (2024-05-22T17:23:15Z)
• A quatum inspired neural network for geometric modeling [14.214656118952178]
  We introduce an innovative equivariant Matrix Product State (MPS)-based message-passing strategy. Our method effectively models complex many-body relationships, suppressing mean-field approximations. It seamlessly replaces the standard message-passing and layer-aggregation modules intrinsic to geometric GNNs.
  arXiv  Detail & Related papers  (2024-01-03T15:59:35Z)
• Equivariant Matrix Function Neural Networks [1.8717045355288808]
  We introduce Matrix Function Neural Networks (MFNs), a novel architecture that parameterizes non-local interactions through analytic matrix equivariant functions. MFNs is able to capture intricate non-local interactions in quantum systems, paving the way to new state-of-the-art force fields.
  arXiv  Detail & Related papers  (2023-10-16T14:17:00Z)
• Set-based Neural Network Encoding Without Weight Tying [91.37161634310819]
  We propose a neural network weight encoding method for network property prediction.<n>Our approach is capable of encoding neural networks in a model zoo of mixed architecture.<n>We introduce two new tasks for neural network property prediction: cross-dataset and cross-architecture.
  arXiv  Detail & Related papers  (2023-05-26T04:34:28Z)
• Neural Functional Transformers [99.98750156515437]
  This paper uses the attention mechanism to define a novel set of permutation equivariant weight-space layers called neural functional Transformers (NFTs) NFTs respect weight-space permutation symmetries while incorporating the advantages of attention, which have exhibited remarkable success across multiple domains. We also leverage NFTs to develop Inr2Array, a novel method for computing permutation invariant representations from the weights of implicit neural representations (INRs)
  arXiv  Detail & Related papers  (2023-05-22T23:38:27Z)
• Permutation Equivariant Neural Functionals [92.0667671999604]
  This work studies the design of neural networks that can process the weights or gradients of other neural networks. We focus on the permutation symmetries that arise in the weights of deep feedforward networks because hidden layer neurons have no inherent order. In our experiments, we find that permutation equivariant neural functionals are effective on a diverse set of tasks.
  arXiv  Detail & Related papers  (2023-02-27T18:52:38Z)
• Equivariant Architectures for Learning in Deep Weight Spaces [54.61765488960555]
  We present a novel network architecture for learning in deep weight spaces. It takes as input a concatenation of weights and biases of a pre-trainedvariant. We show how these layers can be implemented using three basic operations.
  arXiv  Detail & Related papers  (2023-01-30T10:50:33Z)
• Dynamic Graph: Learning Instance-aware Connectivity for Neural Networks [78.65792427542672]
  Dynamic Graph Network (DG-Net) is a complete directed acyclic graph, where the nodes represent convolutional blocks and the edges represent connection paths. Instead of using the same path of the network, DG-Net aggregates features dynamically in each node, which allows the network to have more representation ability.
  arXiv  Detail & Related papers  (2020-10-02T16:50:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.