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

该论文提出了宽度无关的神经网络生成模型 NNiT，它通过图超网络实现权重空间的结构对齐，将架构令牌与权重补丁统一建模，从而在未见过的架构拓扑上实现了卓越的泛化能力。

要点

这篇论文介绍了一种名为 NNiT 的新方法，它的核心目标是：像“生成图片”一样生成神经网络，而且不管这个网络是“胖”是“瘦”（宽度不同），都能直接生成好用的版本。

为了让你轻松理解，我们可以把生成神经网络想象成**“盖房子”**。

1. 以前的难题：盖房子的“死板图纸”

在 NNiT 出现之前，用 AI 生成神经网络（也就是生成控制机器人的“大脑”）面临两个大麻烦：

• 麻烦一：图纸尺寸锁死（宽度不灵活）
  以前的方法就像是在画一张固定尺寸的图纸。如果你训练时只见过“宽 10 米”的房子，那么当你想要盖一个“宽 20 米”的房子时，AI 就懵了。它不知道多出来的空间该怎么填，因为它的“生成逻辑”是绑定在固定尺寸上的。

  • 比喻： 就像你学会了一套乐高积木的拼法，只能拼出 10 块宽的车。如果你想拼 20 块宽的车，以前的 AI 只会把原来的车强行拉长，结果车就散架了。

• 麻烦二：零件顺序混乱（排列对称性）
  神经网络里的神经元（就像盖房子的砖块）其实是可以随意调换顺序的，只要它们之间的连接逻辑对，房子就能住人。但是，对于 AI 来说，砖块 A 在左边和砖块 A 在右边，看起来是完全不同的“数据”。

  • 比喻： 想象你在教 AI 认“人”。如果一个人把头发从左边梳到右边，或者把左边的口袋换到右边，以前的 AI 就会觉得这是两个完全不同的人，完全认不出来了。这导致 AI 很难学会通用的规律。

2. NNiT 的绝招：把“砖块”变成“有纹理的布料”

NNiT 做了两件聪明的事来解决上述问题：

第一步：用“智能织布机”整理砖块（结构化对齐）

作者发现，如果用一种叫 GHN（图超网络） 的工具来生成神经网络的参数，就像是用一台智能织布机。

• 传统的生成方式（SGD）像是在扔砖头，砖块堆在一起乱七八糟，没有规律。
• GHN 则像织布机，它生成的“砖块”（权重）自带纹理和规律。比如，它生成的布料上会有自然的竖条纹或波浪纹。
• 比喻： 以前是给你一堆散乱的沙子，现在 GHN 帮你把沙子压成了有纹理的砖块。因为纹理是连续的，AI 就能看出：“哦，原来这块砖旁边应该接那种纹理的砖”，不管房子多宽，纹理都能接得上。

第二步：把房子切成“小方块”来拼（Patch Tokenization）

这是 NNiT 最核心的创新。它不再把整个神经网络看作一个巨大的、固定的向量，而是把它切成一个个小方块（Patches），就像把一张大图片切成很多小图块。

• 以前： 生成一个 10 米宽的房子，需要生成 100 个特定的数字。
• 现在： 生成一个 10 米宽的房子，就是生成 10 个“小方块”；生成 20 米宽的房子，就是生成 20 个“小方块”。
• 比喻： 就像玩乐高或者像素画。你不需要为每种尺寸重新发明一套积木。你只需要学会怎么拼“一个方块”，然后想盖多宽，就多拼几个方块。
  • 如果用户说：“我要一个更宽的神经网络”，NNiT 就像画家一样，在画布上多画几笔（多生成几个方块），而不是重新画一张新图。

3. 它是怎么工作的？（多模态序列模型）

NNiT 把“房子的设计图”（架构，比如几层、每层多宽）和“房子的砖块”（权重参数）混在一起，当成一串连续的指令来生成。

• 输入： 一串指令，比如 [宽 32 的层] + [砖块纹理 A] + [宽 64 的层] + [砖块纹理 B]...
• 输出： 一个可以直接运行的、完美的神经网络。

它不仅能**“看图说话”（给定架构，生成对应的参数），还能“自由创作”**（不给架构，直接生成一个既好看又好用的新架构和新参数）。

4. 实际效果：机器人真的能干活吗？

作者把这项技术用在了机器人控制（ManiSkill3 任务）上，比如让机器人抓取方块、推箱子。

• 测试场景： 训练时，机器人只见过几种特定宽度的“大脑”。
• 挑战： 测试时，给机器人一个从未见过的、更宽或更窄的“大脑”架构。
• 结果：
  • 以前的方法（Baseline）：完全失效，机器人像喝醉了一样，成功率跌到 0% 或 50%。
  • NNiT： 成功率依然保持在 85% 以上！它成功地把在“窄房子”里学到的经验，迁移到了“宽房子”里。

总结

NNiT 就像是一个拥有“无限缩放能力”的建筑大师。

它不再死记硬背某种固定尺寸的房子图纸，而是学会了**“纹理”和“模块化”**。

1. 它用GHN把混乱的零件整理成有规律的纹理。
2. 它用**切块（Patch）**的方式，让生成过程像拼乐高一样灵活。
3. 无论你想盖多宽、多深的房子，它都能直接生成一套能用的“大脑”，而无需重新训练。

这项技术让 AI 生成神经网络变得更加灵活、通用，为未来在机器人、自动驾驶等领域快速部署定制化 AI 模型打开了新的大门。

技术摘要

1. 研究背景与核心问题

背景：
生成式模型（如扩散模型）在图像、视频和蛋白质设计等领域取得了巨大成功。近期，这一范式被扩展到神经网络参数合成（Parameter Synthesis），即直接生成可工作的神经网络权重，以绕过传统的训练过程。

核心挑战：
现有的神经网络参数生成方法面临两个主要障碍：

1. 排列对称性（Permutation Symmetry）： 神经网络中的神经元顺序是任意的。不同的参数化配置（即权重的排列顺序不同）可以产生完全相同的输入 - 输出函数。这导致相邻权重在空间上缺乏相关性，使得基于空间结构（如卷积或 Patch 处理）的生成模型难以学习。
2. 宽度依赖性（Width-Agnostic Challenge）： 传统方法通常将权重展平为固定维度的向量。一旦网络宽度（Width）发生变化，向量的维度就会改变，导致生成先验（Generative Prior）失效。这使得模型难以泛化到训练集中未见过的网络架构（特别是宽度不同的架构）。

目标：
开发一种能够**宽度无关（Width-Agnostic）**地生成神经网络权重的方法，使其能够零样本（Zero-shot）合成训练集中未见过的任意拓扑结构的网络。

---

2. 方法论 (Methodology)

作者提出了 NNiT (Neural Network Diffusion Transformers)，这是一种多模态扩散 Transformer 框架。其核心思想是将神经网络生成建模为单一序列任务，结合离散架构令牌和连续权重块。

2.1 结构对齐：图超网络 (GHNs) 作为数据生成器

为了解决排列对称性和空间结构缺失的问题，作者利用图超网络 (Graph HyperNetworks, GHNs) 作为权重数据的生成源，并引入关键改进：

• CNN 解码器： 在 GHN 中，作者使用卷积神经网络（CNN）作为解码器，将图节点嵌入映射为权重张量。
• 诱导局部相关性： 与随机梯度下降（SGD）训练产生的无序权重不同，GHN 结合 CNN 解码器会在权重空间中引入显式的局部性偏差（Locality Bias）。这使得生成的权重矩阵在空间上具有高度一致的局部相关性（例如出现垂直带状结构），从而将原本无序的参数空间转化为具有稳定几何结构的“场”。
• 作用： 这种结构对齐消除了排列对称性带来的歧义，使得基于 Patch 的 Tokenization 成为可能。

2.2 多模态序列建模：NNiT 架构

NNiT 将神经网络合成统一为一个序列建模任务：

• 架构 Token（离散）： 将网络架构（层宽序列）表示为离散 Token 序列。
• 权重 Patch（连续）： 利用 GHN 诱导的结构对齐，将权重矩阵视为连续的空间场。权重被切分为 $p \times p$ 的Patch（块），而不是展平为全局向量。
  • 宽度无关性： 增加网络宽度仅意味着生成更多的 Patch，而不改变 Token 化方案或嵌入维度。
• 联合序列： 离散架构 Token 和连续权重 Patch 被拼接成一个统一的序列 $z = [z_a; z_w]$，输入到 Diffusion Transformer (DiT) 中。

2.3 训练策略：混合噪声水平 (MoNL)

为了同时支持联合生成和条件生成，采用了 Mixture of Noise Levels (MoNL) 框架：

• 联合生成模式 (Joint Generation)： 架构和权重同时添加噪声，模型学习联合分布 $p(a, w)$，可从头生成新的架构 - 权重对。
• 条件合成模式 (Conditional Synthesis)： 架构 Token 保持无噪声（$t_a=0$），仅对权重添加噪声，模型学习条件分布 $p(w|a)$，即根据给定架构生成权重。
• 损失函数： 结合了架构和权重的噪声预测误差，以及用于学习协方差的变分下界项。

2.4 部署与合成

在推理阶段，模型根据架构 Token 解码出目标层宽，然后从生成的最大权重张量网格中裁剪出对应的有效子矩阵，组装成可执行的 MLP。

---

3. 主要贡献 (Key Contributions)

1. 证明了 GHN 的结构对齐能力： 发现带有 CNN 解码器的 GHN 能够消除排列对称性，诱导产生具有稳定局部空间相关性的权重分布，为基于 Patch 的生成奠定了基础。
2. 提出了宽度无关的 Patch Tokenization： 首次将神经网络权重视为空间场并进行 Patch 化处理，实现了生成过程与固定矩阵维度的解耦，支持零样本生成未见过的架构拓扑。
3. 构建了 NNiT 多模态扩散 Transformer： 将离散架构搜索与连续参数生成统一在一个序列模型中，实现了联合生成 ($p(a, w)$) 和条件权重合成 ($p(w|a)$)。

---

4. 实验结果 (Results)

实验在 ManiSkill3 机器人控制环境中进行，使用多层感知机（MLP）作为策略网络。

• 结构对齐验证：

  • 对比 GHN 生成与 SGD 训练的权重，GHN 生成的权重在 35 个独立种子下表现出高度一致的局部空间结构（垂直带状），而 SGD 权重则呈现无序噪声。
  • 尽管结构对齐，GHN 生成的策略仍保持了高多样性（L2 距离大，余弦相似度低），未出现模式坍塌。

• 零样本宽度迁移 (Zero-Shot Width Transfer)：

  • 测试设置： 在训练集中未见过的网络宽度（Unseen Topologies）上进行测试。
  • 对比基线： 与 SANE 和 D2NWG 等基线模型对比。
  • 结果：
    • 基线模型： 在未见过的架构上表现极差（成功率降至 0% - 59%），因为它们依赖固定维度的向量表示，无法适应宽度变化。
    • NNiT： 在未见过的架构上保持了 >85% 的高成功率（例如在 [input, 32, 16, 16, 16, output] 上达到 98% 成功率）。这证明了其真正的宽度无关泛化能力。

• 多模态联合生成：

  • NNiT 能够自发地生成完整的网络策略（架构 + 权重），在未见过的拓扑结构上也能达到 90%-100% 的成功率，证明了模型内化了神经网络的结构性逻辑，而非死记硬背训练样本。

---

5. 意义与影响 (Significance)

• 突破架构生成的瓶颈： 解决了神经网络生成中“固定维度”和“排列对称性”两大核心难题，使得生成模型能够像处理图像或视频一样处理可变宽度的神经网络参数。
• 机器人控制与 Sim2Real： 在机器人控制等对权重精度要求极高的领域（微小误差即导致任务失败），NNiT 展示了强大的泛化能力。这为元学习（Meta-learning）、针对特定硬件约束优化网络以及 Sim2Real 部署开辟了新途径。
• 可扩展性： 该方法将网络深度视为时间维度，权重矩阵视为空间特征，类比于视频生成。这使得 NNiT 可以利用视频扩散 Transformer 的优化技术（如线性注意力），未来有望扩展到生成十亿参数级别的基础模型。
• 灵活部署： 单个生成器即可满足用户指定的宽度或计算预算约束，无需为每种架构单独训练策略。

总结：
NNiT 通过引入结构对齐的 GHN 和 Patch 化的权重表示，成功实现了宽度无关的神经网络生成。它不仅能够根据给定架构生成高性能权重，还能联合生成全新的架构与参数组合，在未见过的复杂拓扑上展现出卓越的零样本泛化能力，为神经架构搜索（NAS）和自动化模型生成提供了新的范式。