Pith. sign in

REVIEW 5 major objections 5 minor 2 references

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

T0 review · 5 major / 5 minor · reviewed 2026-08-02 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read A patch-based diffusion transformer can generate functional MLP policies for architectures never seen during training, provided a graph hypernetwork with a convolutional decoder aligns the weight space first.

desk verdict Good idea, clean table, but 'width-agnostic' is an over-claim; the abstract also promises an MNIST result that is not in the paper. read the letter →

arxiv 2603.00180 v2 pith:ML3ZOUAH submitted 2026-02-26 cs.LG cs.AI

Jiwoo Kim , Swarajh Mehta , Hao-Lun Hsu , Hyunwoo Ryu , Yudong Liu , Miroslav Pajic This is my paper
classification cs.LGcs.AI
keywords width-agnosticneuralnetworkgenerationweight-spacediffusionpatchtokenizationgraphhypernetworkspermutationsymmetryzero-shotarchitecturalgeneralizationpolicysynthesismultimodaltransformer
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper is trying to prove that neural-network weight generation does not have to be tied to a fixed layer width. Its move is to replace flattened weight vectors with patch-tokenized weight matrices, and to use a graph hypernetwork with a convolutional decoder to create the local spatial structure that makes patches meaningful. On three simulated robotic manipulation tasks, the resulting model synthesizes policies for held-out topologies with 86–100% success, while prior vectorized baselines fall to 0–59%. If the claim holds, a single generator could produce working networks across a whole family of architectures, and architecture search and weight generation become one joint sampling problem.

What carries the argument

The load-bearing mechanism is patch tokenization of aligned weight tensors: each layer's weight matrix plus bias is padded onto a maximal coordinate grid and cut into non-overlapping p×p patches that are linearly projected into the transformer's embedding space. For this to work, the paper relies on a graph hypernetwork whose CNN decoder acts as a shared, locally structured generator: it maps node embeddings to full weight tensors through convolutional upsampling, producing the consistent vertical banding and local correlations that turn weights into a continuous field. A diffusion transformer backbone with per-modality noise levels then models the joint distribution over discrete architectu

What would settle it

Train the same NNiT pipeline on expert SGD-trained policies that have been explicitly aligned into a common neuron order via permutation matching, then measure zero-shot width transfer; if patch tokenization only works with the graph hypernetwork's particular vertical banding, the success rate on unseen widths will collapse, while success would show that alignment alone, not the generator family, is what matters.

Watch

Extended reading notes

Core claim

The paper claims that graph hypernetworks with a CNN decoder do not just predict weights; they impose a consistent neuron ordering and local spatial correlation on the entire weight distribution, collapsing permutation ambiguity. Once weights are viewed as spatially structured fields, a diffusion transformer can tokenize each weight matrix into p×p patches and learn a joint distribution over architecture tokens and weight patches. This makes width generation analogous to image resolution: unseen wider or narrower layers are handled by generating or cropping patches, not by re-learning the token space. On three simulated robotic tasks, NNiT achieves 99%, 100%, and 86% success on architectures

Load-bearing premise

Everything rests on the claim that weights produced by a single shared generator with local spatial structure are a faithful, aligned stand-in for the space of working policies; if that generator family misses useful policies or its alignment degrades on some topology, the diffusion model learns to imitate a narrow subfamily rather than to synthesize networks.

Editorial extensions

If this is right

  • Width becomes a resolution-like variable: after training, the same model can synthesize weights for wider or narrower layers than it saw, because only the number of patches changes, not the tokenization.
  • Depth becomes a temporal-sequence variable: adding a layer means appending an architecture token and its weight block, so the same sequence model can in principle scale to deeper networks.
  • Architecture search and weight generation collapse into one task: NNiT can sample complete architecture-weight pairs with 90–100% success, rather than searching topologies and then training each one.
  • The generated policies are near-dataset quality: top-10 sampled policies retain 86–100% success, meaning the synthetic weights are directly deployable without test-time optimization.
  • Because synthesis is conditioned on discrete tokens, the same pipeline can condition on any discrete signal, such as a compute budget or task identifier, turning policy generation into a constraint-following problem.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The paper tests widths only up to 64 and depths up to 4; if the width-agnostic claim is genuinely about resolution, the same trained model should extrapolate to 128- or 256-wide layers, and that stress test is a natural next experiment.
  • The alignment story is generator-specific: nothing in the paper measures whether NNiT would work on SGD-trained policies that are post-hoc canonicalized. A head-to-head against weight-matched SGD weights would reveal whether the local banding itself is necessary or whether any consistent ordering suffices.
  • If the graph hypernetwork's structural alignment is what matters, then any deterministic, anchor-conditioned generator with local inductive bias could replace it, decoupling NNiT from the expensive ensemble of 128 hypernetworks used to build the dataset.
  • In practice this reframes architecture search: instead of searching over widths and training each candidate, one could sample functional networks conditioned on a width or compute budget, making hardware-aware deployment a sampling-time operation rather than a training-time one.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

5 major / 5 minor

Summary. The paper proposes NNiT, a diffusion-transformer framework that generates MLP policy weights by tokenizing weight matrices into patches and modeling them jointly with discrete architecture tokens. The central methodological claims are (i) that a Graph HyperNetwork (GHN) with a CNN decoder produces a structurally aligned weight space with stable local correlations, and (ii) that patch-based tokenization makes generation width-agnostic, so synthesizing a wider layer is analogous to increasing image resolution. The authors evaluate NNiT on ManiSkill3 robotics tasks, reporting high success rates on held-out architecture topologies and showing that vectorized baselines (SANE, D2NWG) degrade on unseen structures. The abstract also claims the pipeline generalizes to MNIST classification, but no MNIST experiment appears in the paper.

Significance. If fully validated, the paper would make a useful contribution to neural-network weight generation: it offers a concrete tokenization that avoids fixed-dimension flattening, a joint architecture-weight generative model, and a practical data-generation pipeline based on GHN teachers. The ManiSkill3 evaluation is methodologically strong in several respects: dataset filtering thresholds are disclosed (Appendix A), hyperparameters are tabulated (Appendix C), and the top-10 evaluation protocol is standard for this literature. The diversity metrics (Appendix B, Table 6) and rollout visualizations (Appendix E) are also helpful. However, the load-bearing claims currently outrun the evidence: the MNIST statement is unsupported, the 'width-agnostic' claim is tested only over a discrete vocabulary of widths, and the 'structural alignment' premise rests on qualitative visual inspection rather than a quantitative alignment metric.

major comments (5)
  1. [Abstract] The abstract states: 'the same pipeline also generalizes to MNIST classification beyond the robotic control setting.' I could not find any MNIST experiment, figure, table, or section in the main text or the appendices. This is a scope claim that is not supported by the reported evaluation (§6.2–6.5 covers only ManiSkill3). Either add the MNIST experiment and its results, or remove this sentence from the abstract.
  2. [§5.1.1, §6.2, Table 4] The central claim of width-agnosticism is not tested. The architecture vocabulary is V={16,32,64} (§5.1.1), and all 64 training and 8 held-out topologies are 4-hidden-layer MLPs whose widths are drawn from this same vocabulary (§6.2, Table 4). There is no embedding for a width such as 24 or 128, and §5.1.2's padding to a maximal grid H×W with M as maximum supported width means widths larger than M cannot be cropped out of the generated tensor. Thus Table 2's 'unseen' configurations are unseen combinations of seen widths, not unseen widths. This does not establish the §6.4 claim that patch tokenization 'decouples the generative prior from the global weight space dimensions.' To support the width-agnostic claim, the experiments need to include at least one width outside V (e.g., 24, 48, or 128) and demonstrate that the model can synthesize a functional network at that width, or the claims
  3. [§6.1, Fig. 3] The paper's premise that GHNs 'structurally align the weight space' is evidential only via the visual banding in Fig. 3. Table 1 reports diversity metrics (L2 distance, cosine similarity), but these do not measure alignment across seeds. Since patch-based tokenization and width-agnostic transfer depend on 'consistent local spatial correlations across the population' (§4), the authors should provide a quantitative alignment metric. Concretely, they could compute permutation-invariant correspondence between independently generated weight tensors (e.g., after solving a matching problem, measure residual distance or correlation of coordinates), or measure how consistently functional roles are assigned to spatial positions. The comparison should include SGD-trained networks of matched task performance, as in Table 1.
  4. [§4, §6.1] The alignment claim is demonstrated only on GHN-generated weights. The paper shows that GHN outputs have certain structural properties, not that MLP weight space itself can be aligned. This is a limitation of the current evidence: NNiT is trained exclusively on GHN-generated data, so its success may be specific to the GHN output family. The manuscript should state this limitation explicitly and, ideally, test the pipeline on an alternative aligned weight distribution (e.g., SGD-trained networks after canonicalization) or at least quantify how far GHN weights are from the SGD-trained distribution. Without this, the phrase 'structurally aligned weight space' is a property of the data generator rather than a general property of MLPs.
  5. [§6.3, Table 2] The baseline comparison for zero-shot transfer is not fully specified. The authors state that D2NWG and SANE were 'adapted to this setting' (§6.3), but no details are given about how D2NWG, which uses fixed zero-padded vectorization, was conditioned on unseen architecture topologies during evaluation. Since the comparison is central to the paper's claim that NNiT outperforms baselines on unseen structures, the adaptation protocol should be described (or code released) to rule out an unfair comparison.
minor comments (5)
  1. [Eq. (3)] The total loss in Eq. (3) is written without the balancing terms and the L_vb term is not defined until Appendix C. Please add a forward reference or define L_vb in the main text.
  2. [§5.1.2] Typo: 'we leverage the structural alignment of the structurally aligned weight spaces' — the phrase is redundant. Also, the definition of W = M+p is not motivated; please clarify why the bias column requires an extra column.
  3. [§5.4] The sentence 'The sequence â is decoded into integer widths [n1,...,nS] where n1,...,nS ∈ V' is ambiguous about the input and output widths; clarify that only hidden widths are drawn from V and that input/output sizes are fixed by the task.
  4. [Table 2] In the header, 'Architecture' appears to be a leftover column label; the table has no explicit architecture column. Also, the standard deviations are reported only for returns, not for success rates; please indicate the number of seeds for the success-rate estimates.
  5. [Appendix A] The dataset uses GHNs from HyperPPO (Hegde et al., 2023) but the relationship to the GHN description in §3.2 is not explicit. A short sentence connecting the two would help reproducibility.

Circularity Check

1 steps flagged · score 4.0 of 10

The central 'width-agnostic' claim is self-definitional: all architectures are decoded from a fixed width vocabulary V={16,32,64}, so 'unseen topologies' are unseen combinations of seen widths; functional ManiSkill results remain externally grounded.

  1. self definitional [§5.1.1, §5.4, §6.2, §6.4]
    "We formulate the architecture a∈ V^S as a sequence of discrete tokens, where V is the vocabulary of layer widths ... layer widths sampled from V ∈ {input,16,32,64,output} ... The sequence â is decoded into integer widths [n1,...,nS] where n1,...,nS ∈V ... By decomposing weights into locally consistent patches, NNiT decouples the generative prior from the global weight space dimensions."

    The architecture decoder is restricted to V={16,32,64} (padding to max grid M in §5.1.2), so every 'unseen' test topology in Table 4 is an unseen combination of training widths, not an unseen width. No token or embedding exists for any width outside V, and widths >M cannot be cropped from the padded grid. Therefore the claim that patch tokenization 'decouples the generative prior from the global weight space dimensions' is asserted by construction; the reported success reduces to interpolation over a fixed discrete width vocabulary, not to synthesis across arbitrary/unseen dimensions. The width-agnostic property is untested for widths outside the training set.

full rationale

No other significant circularity is present. The core empirical loop is externally anchored: GHN-generated weights are filtered by ManiSkill3 success, NNiT is trained on that corpus, and the reported 99%/100%/86% zero-shot successes are measured by deploying sampled policies in the simulator, not by construction. The GHN alignment claim is validated by a direct GHN-vs-SGD comparison (Table 1, Figure 3); it is a property of the data generator rather than a redefinition of the target. Baselines are trained on the same expert dataset, so the comparison is fair. The abstract's MNIST claim is unsupported by the body, and the width-vocabulary limitation is a scope concern, but these are not additional circularity. Self-citations (e.g., MoNL, diffusion baselines) are used as background or implementation components and are not load-bearing. The single self-definitional step above concerns the central width-agnostic claim and lowers the score to 4.

Assumptions & free parameters 4 free parameters · 5 assumptions · 0 invented entities

The central method rests on two hand-designed grid choices (patch size p, maximal width M) and a hand-thresholded dataset of GHN-generated weights. No new physical entities are introduced. The main ungated assumption is domain-level: that GHN-generated, filtered weights are a faithful stand-in for the space of functional MLPs, and that the alignment the GHN induces persists on unseen topologies. Everything downstream (NNiT's success) inherits the GHN family's quality and the filtering thresholds.

free parameters (4)
  • Maximal grid width M (padding bound) = 64 (implied by width vocabulary {16,32,64}; stated as 'arbitrary')
    Section 5.1.2 pads all weight blocks to a maximal coordinate grid H x W; this bounds the width-agnostic claim — the model can only assemble widths up to M, and no test uses a width outside the training vocabulary.
  • Patch size p = not stated in paper
    Section 5.1.2 decomposes weight tensors into non-overlapping p x p patches; p controls token count and locality granularity and is never given numerically, making the method's central hyperparameter unspecified.
  • Dataset filtering thresholds (tau_success, tau_return) = 0.9/0.9/0.8 success; 35/35/30 return per task (Table 5)
    Appendix A retains only top-100 GHN policies per architecture passing task-specific thresholds; these hand-chosen thresholds shape the 'expert' distribution NNiT learns and the success rates it can attain.
  • Exclusion rule: width-16 at depths 3-4 = n/a
    Appendix A excludes architectures with width 16 at depths 3 and 4 'due to insufficient capacity'; a post-hoc data selection that trims the search space before the train/test split.
assumptions (5)
  • standard math Permutation symmetry of MLP losses and functional invariance under Eq. 1
    Section 3.1 formalizes that the input-output function is invariant to hidden-unit permutations; this is background theory from Hecht-Nielsen 1990 and Ainsworth et al. 2023.
  • standard math Gaussian forward diffusion and noise-prediction objective (Eq. 3, 5)
    Standard DDPM assumptions; the model predicts noise and a covariance via a variational term, following Ho et al. 2020.
  • domain assumption GHN weight distribution is representative of the space of functional MLP policies
    Entire evaluation flows from GHN-generated weights (Appendix A); no experiments show NNiT transfers to SGD-trained networks or to policies trained outside the GHN family. This is the weakest load-bearing assumption.
  • domain assumption The GHN-CNN-decoder locality bias persists for held-out topologies
    Section 4 and 6.1 argue the CNN decoder induces consistent local correlation; the untested extension is that this structural alignment holds for unseen width combinations, not just the training topologies.
  • domain assumption Top-10-of-100 policy selection with 50-episode evaluation approximates deployment utility
    Section 6.3 reports performance of the top-10 of 100 generated samples, simulating offline validation; deployment of the 'generated network' therefore includes an expensive selection step.

how reviews work

0 comments
Cite this review

Pith. "Pith review of NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces." pith.science (2026). https://pith.science/paper/ML3ZOUAH

@misc{pith2026260300180,
  author       = {Pith},
  title        = {Pith review of: NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/ML3ZOUAH}},
  note         = {Machine review of arXiv:2603.00180}
}
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 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.

Figures

Figures reproduced from arXiv: 2603.00180 by the authors.

Figure 1
Figure 1. Width-Agnostic Synthesis via Multimodal Tokenization. Unlike previous models, NNiT decouples functional logic from fixed matrix dimensions, allowing the zero-shot synthesis of optimal weights for architectural topologies entirely unseen during training. sequence modeling task. NNiT tokenizes aligned weight tensors into p × p patches, replacing global vectors with spatially correlated tokens. This representation make… view at source ↗
Figure 2
Figure 2. NNiT Framework Overview. Left: Unified Generative Architecture. We formulate neural synthesis as a multimodal sequence task. Discrete architecture tokens (orange) and continuous weight matrices (blue) are unified into a single sequence, with weights processed as spatially correlated patches. A Diffusion Transformer (DiT) models the joint distribution using per-modality timestep conditioning (µa, Σa, µw, Σw) via the … view at source ↗
Figure 3
Figure 3. Visualizing Structural Alignment and Induced Geometry. Comparison of weight magnitude profiles across 35 independent seeds. Top (GHN): The consistent alignment across the seeds demonstrates that the GHN 1) successfully spatially aligns the weight spaces, effectively resolving the permutation ambiguity inherent in neural networks. Furthermore, the visible structural banding indicates that the GHN 2) imposes meaningfu… view at source ↗
Figures from the paper (6 more)
Figure 4
Figure 4. Figure 4: Visualization of Topological Anchoring. Heatmaps of neuron-wise weight magnitudes for 3 selected architectures across 100 filtered seeds. The vertical banding visually confirms the induction of spatial correlation, validating the premise that these weights can be treat…
Figure 5
Figure 5. Figure 5: Dataset Diversity Analysis. Histograms of pairwise L2 distances and Cosine Similarities across all three environments. The consistently high L2 distances and low cosine similarities confirm that the structural alignment imposed by the GHN does not result in mode collap…
Figure 6
Figure 6. Figure 6: Parametric Diversity Analysis. Pairwise Euclidean distance (left) and cosine similarity (right) matrices for the top 10 generated policies. The dominance of high distances and low cosine similarities confirms that the generative model preserves variance within the para…
Figure 7
Figure 7. Figure 7: PickCube-v1 Policy Rollout. A sequential visualization of policies generated by NNiT across diverse topologies. Left (Y-axis): Target architecture configurations denoted by their hidden layer widths (e.g., 32-16-64-32). Top (X-axis): Temporal snapshots of the rollout i…
Figure 8
Figure 8. Figure 8: PushCube-v1 Policy Rollout. A sequential visualization of policies generated by NNiT across diverse topologies. Left (Y-axis): Target architecture configurations denoted by their hidden layer widths (e.g., 32-64-32-64). Top (X-axis): Temporal snapshots of the rollout i…
Figure 9
Figure 9. Figure 9: StackCubeEasy-v1 Policy Rollout. A sequential visualization of policies generated by NNiT across diverse topologies. Left (Y-axis): Target architecture configurations denoted by their hidden layer widths (e.g., 64-64-64-32). Top (X-axis): Temporal snapshots of the roll…

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

2 extracted references · 2 linked inside Pith

  1. [2023]

    URL https://openreview.net/forum? id=heTTJfkTQC. Kim, M., Pertsch, K., Karamcheti, S., Xiao, T., Balakr- ishna, A., Nair, S., Rafailov, R., Foster, E., Lam, G., Sanketi, P., Vuong, Q., Kollar, T., Burchfiel, B., Tedrake, R., Sadigh, D., Levine, S., Liang, P., and Finn, C. Open- vla: An open-source vision-language-action model.arXiv preprint arXiv:2406.092...

    arXiv 2021

  2. [2025]

    Rinaldi, F., Capitani, G., Bonicelli, L., Crisostomi, D., Bolelli, F., Ficarra, E., Rodol `a, E., Calderara, S., and Porrello, A

    URL https://openreview.net/forum? id=o3pJU5QCtv. Rinaldi, F., Capitani, G., Bonicelli, L., Crisostomi, D., Bolelli, F., Ficarra, E., Rodol `a, E., Calderara, S., and Porrello, A. Update your transformer to the latest re- lease: Re-basin of task vectors, 2025. URL https: //arxiv.org/abs/2505.22697. Ruan, L., Ma, Y ., Yang, H., He, H., Liu, B., Fu, J., Yuan...

    arXiv 2025

Pith tools

Reviewed August 2, 2026 · model on record in the stance chip above.