On the Expressive Power of Permutation-Equivariant Weight-Space Networks¶

会议: ICML 2026
arXiv: 2602.01083
代码: https://github.com/dayanadir/capacity_increase_inr_editing_experiment
领域: 权重空间学习 / 表达力理论 / 等变神经网络
关键词: weight-space learning, permutation equivariance, universality, INR editing, OCE

一句话总结¶

本文为操作在 MLP 权重上的置换等变 weight-space 网络（DWS / NFN / GMN / NG-GNN 等）建立了首个系统的表达力理论，证明这些架构在表达力上几乎完全等价，并在"general position"假设下对四种逼近场景（函数空间泛函/算子、置换不变泛函、置换等变算子）给出了普适性刻画；由理论得出的简单修改 OCE（输出端 ensemble 多个 MLP）在 INR 编辑基准上相对 SOTA 提升 34%。

研究背景与动机¶

领域现状：weight-space learning 把训练好的神经网络当成"结构化数据"，用一个 meta-network 直接吃另一个 MLP 的参数 \(v=(W_1,b_1,\dots,W_L,b_L)\)，做精度预测、INR 编辑、meta-optimization 等下游任务。由于 MLP 隐藏层神经元置换 \(\tau\) 满足 \(f_{\rho(\tau)v}=f_v\)（函数不变），主流 SOTA 架构（DWS / NP-NFN / HNP-NFN / GMN / NG-GNN / NFT）均显式构造为对 \(G_A=S_{d_1}\times\cdots\times S_{d_{L-1}}\) 置换等变的。

现有痛点：(1) 这些架构看起来差别很大（DWS 是手工对齐的等变线性层，GMN/NG-GNN 是把网络当成图跑 GNN，NFT 是 attention），但社区不清楚它们到底谁强谁弱；(2) 对称性约束本身会削弱表达力，但既有理论仅给出零散的 forward-pass 模拟结论 (Navon et al., 2023; Lim et al., 2023; Kalogeropoulos et al., 2024)，没有统一的 universality 刻画；(3) 一些自然的目标，例如 INR 的"放大/zoom-out"算子，输出函数的复杂度可能超过输入 MLP 的容量上限，从理论上看根本不可能用同架构输出去逼近——但这件事此前没有被明确指出。

核心矛盾：weight-space 的逼近问题天然横跨两层语义——一边是参数空间 \(\mathcal V\)，一边是它们实现的函数空间 \(\mathcal F\)。"等变 weight-to-weight 映射"和"函数空间算子"是不同的目标类型，必须分开讨论；而对称性约束在某些设定下足够，在另一些设定下不够（即非普适），需要精确刻画"边界"。

本文目标：(a) 把所有主流置换等变 weight-space 架构在表达力意义下放进同一个等价类；(b) 把"逼近目标"系统分成 4 类并给出 universality / non-universality 的完整图谱；(c) 把理论翻译成可落地的架构改动。

切入角度：作者注意到，几何深度学习里其它对称域（图、点云）常用"在退化子集 \(\mathcal E\) 外才普适"的 general position（GP）范式 (Maron et al., 2020; Finkelshtein et al., 2025)；权重空间中天然的退化集合就是"某一隐藏层存在两个相同偏置 \(b_i=b_j\)"——这是 Lebesgue 测度零的子集，几乎所有训练得到的 MLP 都落在 GP 上。以 GP 为抓手，等变性带来的退化就能被分离出来。

核心 idea：先证所有主流架构表达力等价 → 把分析对象统一成"通用置换等变 weight-space 网络" → 在 4 个逼近场景里逐个判定 universality；同时由"函数空间算子在固定架构下不普适"这一不可能性结论，反推出一个简单破解方法：让输出比输入更大（OCE 直接输出 \(k\) 个 MLP 取平均）。

方法详解¶

本文是纯理论 + 一个理论驱动的简单实验性架构改动。"方法"应理解为定理框架与证明骨架，而不是新模型。

整体框架¶

把 weight-space 问题组织成"目标类型 × 输入是否 GP"的二维网格：

目标类型（Definition 4.1）：
1. 函数空间泛函 \(\Psi:\mathcal C(X,\mathbb R^{d_L})\to\mathbb R^n\)，例如精度预测、INR 分类；
2. 置换不变泛函 \(\Psi:\mathcal V_A\to\mathbb R^n\)，例如权重 \(\ell_2\) 范数、loss landscape 曲率；
3. 函数空间算子 \(\Psi:\mathcal C\to\mathcal C\)，例如 INR 编辑、domain adaptation；
4. 置换等变算子 \(\Psi:\mathcal V_A\to\mathcal V_A\)，例如 pruning mask 预测、meta-optimization 的梯度预测。
输入域：整个 \(\mathcal V\) vs. GP 子集 \(\mathcal V\setminus\mathcal E\)，其中 \(\mathcal E_A=\{v\mid\exists\ell\in[L-1],\,i\ne j,\,(b_\ell)_i=(b_\ell)_j\}\)。

逼近以 \(L^2\)（泛函/等变算子）或 \(L^\infty\)（函数空间算子）度量。对函数空间算子，逼近以 \(\sup_v\|\Psi(f_v)-f_{\Phi(v)}\|_\infty<\epsilon\) 度量，即"等变的权重到权重映射"在被 realization map \(R(v)=f_v\) 拉到函数空间后近似目标算子。

关键设计¶

架构等价定理（Theorem 5.2 + Proposition 5.3）：
- 功能：把 6 个主流架构 \(\Pi=\{\text{DWS, NP-NFN, HNP-NFN, GMN, NG-GNN, NFT}\}\) 折叠成一个表达力等价类。
- 核心思路：对任意紧集 \(K\subseteq\mathcal V\) 和任意 \(\pi,\pi'\in\Pi\setminus\{\text{NFT}\}\)，\(\mathcal N^\pi_{\text{inv}}(K)=\mathcal N^{\pi'}_{\text{inv}}(K)\) 且 \(\mathcal N^\pi_{\text{equi}}(K)=\mathcal N^{\pi'}_{\text{equi}}(K)\)。证明方式是用一个架构的基本层"模拟"另一个架构的基本层（mutual approximation）。NFT 因为 attention 机制非标准，在全空间上不等价（Proposition 5.3 给出反例 \(K\)），但在 GP 子集 \(K\subset\mathcal V\setminus\mathcal E\) 上仍然等价。
- 设计动机：把"我该用哪个架构"从架构选型问题降级为工程问题，并允许后续所有定理只针对"通用置换等变 weight-space 网络"陈述一次。
GP 假设下的统一普适性（Theorems 6.1 / 6.3 / 7.2 / 7.4）：
- 功能：在 4 个目标类别里分别给出 universality 的判定（含上界和下界）。
- 核心思路：
  - 函数空间泛函（Thm 6.1）：在 全空间 \(K\subseteq\mathcal V\) 上就普适。证明用 DWS 可以模拟 MLP forward pass (Navon et al., 2023) → 推出"DWS 不可区分 \(v,v'\) ⇒ \(f_v=f_{v'}\)"的 separation 性质 → 套 Pacini et al. (2025b) 的 separation-to-approximation 定理收尾。
  - 置换不变泛函（Prop 6.2 + Thm 6.3）：全空间下 不普适（构造 \(v,v'\) 使 \(W_2,W_2'\) rank 不同但 1-WL 不可区分，见 Figure 3），但在 GP 上普适。Thm 6.3 的关键构造是连续 canonization 映射 \(\operatorname{canon}:K\to\mathcal V\)：因 \(K\cap\mathcal E=\varnothing\)，每层偏置 \(b_\ell\) 元素两两不同，可用 \(\operatorname{argsort}(b_\ell)\) 作为该层置换，得到的轨道代表既唯一又连续；DWS 含 DeepSets 原语，对 ranking 普适 (Segol & Lipman, 2019)，再接 MLP head 即可。
  - 函数空间算子（Prop 7.1 + Thm 7.2）：在 固定 ReLU 架构 上 不普适——ReLU MLP 的 linear region 数量被 (Montúfar et al., 2014) bound 住，因此放大/zoom-out 这类增加几何复杂度的算子不可被同架构等变映射逼近；但只要允许"输出架构 \(A\) 足够大"，则普适。证明三步：用 partition of unity 把 \(\Psi(f_v)\) 写成 \(M\) 个参考 MLP 的连续凸组合 \(\Rightarrow\) 用 canonization 把它转成置换等变 \(\Rightarrow\) 套 Thm 7.4 收尾。
  - 置换等变算子（Prop 7.3 + Thm 7.4）：与不变情形对偶——通过 broadcasting，invariant 普适性蕴含 equivariant 普适性。Thm 7.4 用"广播 canonization" \(\widetilde{\operatorname{canon}}(v)\)（把 \(\operatorname{Flat}(\operatorname{canon}(v))\) 拼到每个权重/偏置 entry 上）+ pointwise MLP 给出构造。
- 设计动机：四个 cell 全部判定后，得到 Figure 1 的"表达力地图"——既告诉实践者"什么场景现有架构够用"，也精确标出"什么场景需要新架构"。GP 假设把"理论非普适"和"实际几乎总能逼近"两边都说清楚，避免悲观结论被滥用。
OCE：Output Capacity Expansion（Section 8）：
- 功能：把 Thm 7.2 的"输出架构要更大"做成最小化、即插即用的改动。
- 核心思路：在任何 weight-space 网络的最后特征维加一个 \(k>1\) 的维度，把输出张量解读成 \(k\) 个并行 MLP 的参数，最终预测是 \(k\) 个 MLP 输出的平均。代码上几乎只动一行；模型参数量也基本不变（共享了 backbone，只是输出 head 通道扩 \(k\) 倍）。
- 设计动机：Prop 7.1 的不可能性根源是"输出 MLP 容量 \(=\) 输入 MLP 容量"导致 linear region 数不够；ensemble 后等效 ReLU region 数翻 \(k\) 倍，绕开容量瓶颈，且天然保持置换等变（每个分支独立等变）。这是用理论直接 buy 来的工程提升。

损失函数 / 训练策略¶

理论部分无训练目标；OCE 实验沿用 INR dilation benchmark (Zhou et al., 2023a) 的标准 MSE 监督。

实验关键数据¶

实验只有一个：MNIST INR dilation benchmark，用于验证 OCE 带来的 SOTA 增益，从而间接验证 Thm 7.2 的实践价值。

主实验¶

方法	参考	MSE (\(\times 10^{-2}\), ↓)
NFT	Zhou et al. 2023b	5.10 ± 0.04
NP-NFN	Kofinas et al. 2024	2.55 ± 0.00
NG-GNN-64	Kofinas et al. 2024	2.06 ± 0.01
ScaleGMN-B	Kalogeropoulos et al. 2024	1.89 ± 0.00
NG-T-64	Kofinas et al. 2024	1.75 ± 0.01
ScaleGMN + GradMetaNet++	Gelberg et al. 2026	1.60 ± 0.01
DWS (k=1, baseline)	Gelberg et al. 2026	2.29 ± 0.01
GMN (k=1, baseline)	Gelberg et al. 2026	1.96 ± 0.02
DWS + OCE (k=8)	This paper	1.36 ± 0.03
GMN + OCE (k=8)	This paper	1.06 ± 0.13

GMN+OCE 相对此前 SOTA（ScaleGMN+GradMetaNet++ 的 1.60）下降 34%；DWS 和 GMN 自身相对 \(k=1\) 分别下降 41% 和 46%。

消融实验¶

附录 Table 2 的趋势（正文转述）：

配置	关键现象	说明
DWS, \(k=1\to 8\)	MSE 下降 ~41%	不增加额外参数（共享 backbone）
GMN, \(k=1\to 8\)	MSE 下降 ~46%	验证 Thm 7.2 的"扩输出架构"指引
对照 baseline	大量利用 gradient / probe 等额外信号	OCE 不用任何额外信号即超过

关键发现¶

理论 → 实验闭环：性能瓶颈定位为"输出表示容量不足"而非"backbone 弱"，这是 Prop 7.1 直接预言并由 OCE 验证的。
OCE 是 weight-space 学习的 free lunch：参数量几乎不变、对 DWS / GMN 都通用、不需要额外监督信号，却能拉开数十百分点的 MSE 差距，对未来基线设置有强示范作用。
NFT 出乎意料地排名靠后（5.10 vs DWS+OCE 1.36），暗示 attention 在 weight-space 的优势远不如在序列建模——这与 Prop 5.3 中"NFT 在全空间下与其它架构不等价"的理论观察相呼应。

亮点与洞察¶

把一堆看起来差异巨大的架构折叠成一个等价类，是这篇 paper 最高密度的洞察：以后做 weight-space 研究时，"选 DWS 还是 GMN 还是 NG-GNN"几乎只剩工程偏好；论文里只需选一个最方便分析的（这里是 DWS）来证一切。
GP 假设的双重用法：既用来把"反例"和"普适"切开（Prop 6.2 vs Thm 6.3 配对出现），又用来把 NFT 拉回等价类（Prop 5.3）。"在测度零退化集外讨论普适性"这套方法论可以迁移到其它对称域（如带共享参数的 transformer、scale-equivariant 网络）。
连续 canonization 是一把万能钥匙：因为 GP 下 \(\operatorname{argsort}(b_\ell)\) 唯一且局部常数，连续 canonization 自然存在；这把"等变普适性"问题归约为"DeepSets 普适性"，证明骨架因此异常干净，几乎所有 4 个 cell 都共用这一步。
Prop 7.1 与 OCE 的工程意义：以前社区觉得"INR 编辑做不上去"是模型不够强；本文指出根本原因是输出 MLP 容量被锁死了。这个观察直接催生了 OCE 这种"几乎零成本"的改动——可类比到任何"输入小网络→输出小网络"的 meta-learning 设定。
与过参数化的隐喻：作者明确把"扩大输出架构"与"过参数化更易优化、泛化更好" (Du et al., 2019; Belkin et al., 2019) 类比，提示 weight-space 学习的下一个突破口可能在 input-smaller-than-output 的非对称设计上。

局限性 / 可改进方向¶

只覆盖 MLP 权重空间，不含 transformer / conv 权重（虽然作者在附录 H 给了 transformer 的延伸 sketch），CNN 上的卷积+池化对称群尚未处理。
明确排除了 scale-equivariant 架构（如 ScaleGMN, Kalogeropoulos et al. 2024; Tran et al. 2024），但实验 baseline 里 ScaleGMN 表现最强，理论与实践之间有一段缺口。
理论只谈表达力，不谈优化与泛化——作者自己也强调这点；某个映射"可被某网络逼近"不等于"梯度下降能找到它"，在 weight-space 这种 loss landscape 极不规则的领域，gap 可能很大。
不可能性结论依赖 ReLU（Prop 7.1 用 linear region 计数），其它激活下需要重新构造退化函数族；论文承认"相信能推广"但未给完整证明。
OCE 的 ensemble 数 \(k\) 是超参，且只在 INR dilation 这一个 benchmark 上验证，更多 function-space operator 任务（domain adaptation、NeRF 编辑）尚无实证。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次为整个 weight-space 网络家族给出统一的表达力刻画，且把架构等价性、四象限普适性、不可能性、工程修复闭环写在一篇里。
实验充分度: ⭐⭐⭐ 理论占绝对主体，实证只有一个 benchmark；好在结果硬（34% SOTA gain）且直接呼应理论。
写作质量: ⭐⭐⭐⭐⭐ Figure 1 的"表达力地图"一图胜千言，定理-反例-定理的节奏清晰，证明 sketch 自洽可读。
价值: ⭐⭐⭐⭐⭐ 既给后续 weight-space 论文省去"选架构"这一节，也提示了 input-smaller-than-output 的新设计方向；OCE 本身就是即插即用的实用 trick。