Weight-Space Mixture-of-Experts for Implicit Neural Representation Classification¶
会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/stasiek-j/HMoE-MWT
领域: 可解释性
关键词: 隐式神经表示 (INRs), 权重空间学习 (Weight-Space Learning), 混合专家系统 (Mixture-of-Experts), 层次化路由, 权重可解释性 (Grad-CAG)
一句话总结¶
针对隐式神经表示(INR)权重空间分类中参数维度极高且层级频域功能异构的难题,本文提出结合元学习的层次化混合专家(HMoE)Transformer 架构,通过层级与 Token 双阶段稀疏条件计算刷新多项基准 SOTA,并首创基于梯度的权重空间可解释性归因与剪枝方法(Grad-CAG)揭示出判别性子电路的存在。
研究背景与动机¶
隐式神经表示(INRs)利用坐标多层感知机(MLP)的权重将多维连续信号隐式编码,为统一处理图像、音频与 3D 几何等异构多模态数据,以及在联邦学习与隐私保护医疗分析中避免直接共享原始像素提供了全新范式。然而,直接在权重空间中执行下游分类等模式识别任务面临着严峻的技术瓶颈:MLP 内在的神经元排列与缩放对称性极易导致表示歧义,网络参数的极高维度对分类器的容量与压缩效率提出了苛刻要求,且为每个样本独立拟合 INR 会带来高昂的计算开销。
尽管 Meta Weight Transformer (MWT) 引入 MAML 框架通过元学习初始化与端到端联合训练缓解了初始化异构并组织了权重流形,但现有的权重分类器(如密集 Transformer)将所有权重神经元视作可交换的同质 Token,对每个 Token 施加完全共享的前馈网络变换。这种统一处理方式彻底忽视了 INR 权重的本质先验——即不同的网络层严格对应信号的不同频率分量,且神经元在网络结构中承担着显著异构的功能分工。密集 Transformer 既无法自适应分配计算资源,也无法提供权重空间的决策可解释性。
混合专家系统(MoE)天然具备按需激活特定专家子模块的条件计算能力,恰好与 INR 权重空间中层级功能异构与子电路专门化的物理属性高度契合。核心 idea:构建与 INR 网络拓扑解耦对齐的层次化混合专家(HMoE)Transformer,通过层级感知路由与跨层注入的细粒度 Token 路由两阶段协同处理权重特征,并利用一阶敏感度权重归因(Grad-CAG)首次打通权重空间分类的透明度与子电路定位。
方法详解¶
整体框架¶
整体框架由端到端元学习训练流线与层次化权重分类网络组成。在输入端,利用元学习机制维护一个全局共享的基础初始化参数 \(\theta\) 与逐参数学习率调度表 \(\alpha\);对于输入样本,在内循环中利用均方误差重建损失迭代更新 \(k\) 步生成样本专属参数 \(\phi\)。分类器将隐层神经元的传入权重拼接入偏置形成维度为 \(D = n + 1\) 的 Token,总计 \(M = nL\) 个神经元 Token。分类器以残差权重 \(\lambda(\phi - \theta)\) 为输入(\(\lambda = 500\) 放大细微扰动),送入堆叠了 10 个 HMoE 模块的 Transformer 中,每个模块依次执行层级路由与细粒度 Token 路由,最终经均值池化输出类别预测。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入样本与坐标网格"] --> B["元学习内循环自适应<br/>k步优化生成特定INR权重 ϕ"]
B --> C["权重残差Token化<br/>λ(ϕ - θ) 按层聚合为特征矩阵"]
C --> D["层级化混合专家路由<br/>层级均值池化引导 Top-1 层专家分派"]
D --> E["跨层级门控投影协同<br/>将层级路由信号线性投影并偏置 Token 门控"]
E --> F["细粒度 Token 级专家运算<br/>Top-1 专家稀疏计算输出更新表征"]
F --> G["分类与负载均衡联合反向传播<br/>L_cls 反传塑形元INR + L_balance 规避路由塌缩"]
G --> H["权重空间梯度激活图 Grad-CAG<br/>计算一阶重要性并指导剪枝归因分析"]
关键设计¶
1. 层级化混合专家路由:分层解耦层级全局结构与神经元细粒度特征
针对密集前馈网络将不同层神经元一视同仁、无法捕捉频域解耦结构的缺陷,模型将权重序列 \(X \in \mathbb{R}^{M \times D}\) 显式重组为按层划分的子集合 \((X_1, \dots, X_L)\),其中每个子块 \(X_\ell \in \mathbb{R}^{n \times D}\) 聚合了第 \(\ell\) 层全部 \(n\) 个神经元。为了捕获整层的全局频率特性,模块对层内所有 Token 计算均值向量 \(v_\ell = \frac{1}{n} \sum_{i=1}^n X_\ell[i]\),并利用线性投影矩阵 \(W_g\) 与 Softmax 计算该层的路由权重向量: $\(g_\ell = \text{softmax}(W_g v_\ell)\)$ 随后利用 \(\text{TopK}(g_\ell, k)\) 机制挑选前 \(k\) 个层级专家 \(E_j\),整层内的所有神经元 Token 统一由这些被选中的专家并行处理并根据门控权重加权求和,生成层级输出 \(X_\ell'\)。最后将各层输出按原有拓扑顺序重新拼接,既严格保留了层级的宏观结构先验,又极大地降低了专家调度的决策扰动。
2. 跨层级门控投影协同:将层级先验平滑注入 Token 级专家分配
在完成了粗粒度的层级专家处理后,单一层内不同的神经元在子功能表达上仍存在微观差异。如果直接在 Token 级做独立路由,容易割裂该 Token 与所在层的宏观归属。为此,模型在第二阶段设计了跨层级门控投影矩阵 \(B \in \mathbb{R}^{E_L \times E_T}\),将第一阶段生成的层级路由概率分布 \(g_\ell\) 映射至 Token 级专家空间中作为偏置项。对于属于第 \(\ell\) 层的任意单个神经元 Token \(x \in X_\ell\),其 Token 级路由打分向量 \(s\) 计算为: $\(s = \text{softmax}(W_s x + g_\ell B)\)$ 各个 Token 根据打分向量 \(s\) 独立挑选 \(\text{TopK}(s, k)\) 个 Token 专家执行细粒度前馈变换:\(x' = \sum_{j \in \text{TopK}(s, k)} s_j E_j(x)\)。这种由粗到细的层次化联动机制,使得 Token 级路由决策充分吸收了宏观层级约束,实现了频域宏观上下文与神经元局域功能专精的统一。
3. 元学习任务协同驱动:联合反向传播塑造判别性权重流形
为了使拟合图像信号的 INR 权重空间具备良好的线性可分性,系统避免了孤立训练 INR 的传统割裂做法,而是采用端到端 MAML 协同训练框架。元网络(SIREN)维护共享初始化参数 \(\theta\) 与针对 \(k\) 步更新分别定义的学习率向量 \(\alpha \in \mathbb{R}^{k \times |\theta|}\)。在内循环中,模型以坐标重建 MSE 损失快速自适应;在完成外循环预测后,分类交叉熵损失 \(L_{cls}\) 与专家负载均衡损失联合回传至 \(\theta\) 和 \(\alpha\)。外循环梯度直接干预内循环权重的演化轨迹,迫使各样本在权重空间中不仅能够精确重建高保真像素,而且能自主形成彼此聚类、易于判别的连续权重子流形。
4. 权重空间梯度激活图(Grad-CAG):基于一阶敏感度定位判别子电路
传统 Grad-CAM 依赖像素网格或特征图空间,完全无法解释在离散权重空间中决策形成的机理。为此,本文首创了针对权重空间的一阶梯度归因方法 Grad-CAG。对于分类器关于目标类别 \(i\) 的预测输出 \(c_i(\phi)\),每个 INR 标量权重 \(\phi_j\) 的重要性得分 \(S_j^{(i)}\) 定义为: $\(S_j^{(i)} = \left| \phi_j \frac{\partial c_i(\phi)}{\partial \phi_j} \right|\)$ 该打分衡量了特定权重微小扰动对目标类置信度的边际影响。通过将所有权重的打分拉平并利用 UMAP 进行低维投影,可以直观观察权重空间的流形结构;更重要的是,利用该打分指导结构化权重剪枝,能够精准识别出对下游分类至关重要的特定功能子网络(Circuits)。将剪枝后的 INR 重新渲染并与原图做残差后处理,可以直接反解出空间维度的类激活显著图。
损失函数 / 训练策略¶
系统整体目标函数由坐标重建损失 \(L_{rec}\)、分类交叉熵损失 \(L_{cls}\) 以及两级 MoE 专家的负载均衡损失 \(L_{balance}\) 加权构成。为防止专家路由塌缩,对层级 MoE 与 Token 级 MoE 分别计算门控选择频率与平均门控概率的内积: $\(\mathcal{L}_{\text{balance}}^{\text{layer}} = E_L \sum_{e=1}^{E_L} s_e P_e, \quad \mathcal{L}_{\text{balance}}^{\text{token}} = E_T \sum_{j=1}^{E_T} \tilde{s}_j \tilde{P}_j\)$ 单块平衡损失取二者均值,全模型总损失按权重超参组合: $\(\mathcal{L}_{\text{total}} = \mathcal{L}_{rec} + w_{task} \mathcal{L}_{cls} + w_{balance} \mathcal{L}_{balance}\)$ 默认设置 \(w_{task} = 0.01\),\(w_{balance} = 0.1\)。层级与 Token 级阶段均采用 \(k = 1\)(单专家激活)且各配置 4 个专家(\(E_L = 4, E_T = 4\)),使得推理时的理论算力开销(FLOPs)严格与密集 Transformer 保持同一基线。分类器采用 AdamW 优化器(学习率 0.0001,权重衰减 0.0001),训练周期为 40 epochs。
实验关键数据¶
主实验¶
论文在低分辨率基准(MNIST、Fashion-MNIST、CIFAR-10)与高分辨率挑战集(Imagenette、ImageNet-1K)上均进行了全面评测。在所有对比实验中,HMoE 系列模型均大幅超越了以往基于图神经网络(ScaleGMN、NG-GNN)以及密集注意力(MWT)的方法。
| 数据集 | 指标 | HMoE-MWT (本文) | MWT-L* (基线) | 性能提升 |
|---|---|---|---|---|
| MNIST | 分类准确率 (%) | 99.06 ± 0.15 | 98.91 ± 0.11 | +0.15% |
| Fashion-MNIST | 分类准确率 (%) | 90.72 ± 0.25 | 90.35 ± 0.20 | +0.37% |
| CIFAR-10 | 分类准确率 (%) | 65.01 ± 0.63 | 58.90 ± 0.35 | +6.11% |
| Imagenette | 分类准确率 (%) | 62.52 ± 0.60 | 60.62 ± 0.37 (MWT-L) | +1.90% |
| ImageNet-1K | Top-1 准确率 (%) | 26.73 (HMoE-MWT-L) | 24.11 (MWT-L) | +2.62% |
在 Imagenette 的严格等参数量控制对比(Parameter-Matched)中,当参数量均固定为 1.1M 时,压缩版的 [email protected] 达到 61.47 ± 0.32% 准确率(FLOPs 仅 1.9G),显著优于标准 MWT 的 56.78 ± 0.27%(FLOPs 2.3G);当扩大到 11M 参数量时,HMoE-MWT 取得 62.52 ± 0.60%(FLOPs 8.8G),亦优于扩展版 MWT-PM@11M 的 59.92 ± 0.35%(FLOPs 15.6G)。
消融实验¶
在 CIFAR-10 上对 HMoE-MWT 的专家路由层级与专家容量配置进行消融验证(分类器均采用默认 10 层架构):
| 配置与路由层级 | TopK(Layer) | # Layer Exp | TopK(Token) | # Token Exp | 准确率 (%) | 重建 PSNR (dB) | 参数量 | 说明 |
|---|---|---|---|---|---|---|---|---|
| 仅 Layer MoE | 1 | 4 | - | - | 63.21 | 33.86 | 11M | 缺少 Token 级微观细化 |
| 仅 Token MoE | - | - | 1 | 4 | 62.57 | 30.70 | 11M | 缺失宏观层级结构先验 |
| 仅 Token MoE (多活) | - | - | 2 | 4 | 62.54 | 31.14 | 11M | 激活更多专家未见收益 |
| 仅 Token MoE (密集全开) | - | - | 4 | 4 | 62.21 | 30.68 | 11M | 退化为稠密混合计算 |
| 完整 HMoE (默认配置) | 1 | 4 | 1 | 4 | 65.01 | 31.71 | 11M | 双阶段分层协同最优 |
| 扩增层级容量 | 1 | 8 | 1 | 4 | 64.03 | 31.86 | 17M | 扩容未带来增益 |
| 扩增 Token 容量 | 1 | 4 | 1 | 8 | 64.53 | 32.27 | 17M | 仅略微提升重建 PSNR |
| 双阶段全量扩容 | 1 | 8 | 1 | 8 | 63.64 | 30.61 | 22M | 过拟合并略有性能退化 |
关键发现¶
- 分层解耦路由至关重要:单独使用 Token-wise MoE 准确率仅为 62.57%,而加入 Layer-wise 粗粒度路由后性能跃升至 65.01%(绝对提升 2.44%),证明尊重 INR 层级与频域结构是权重学习的先决条件。
- 稀疏路由优于稠密分配:在 Token 阶段将 Top-k 从 1 增加到 2 或 4,准确率反而从 62.57% 下滑至 62.21%,且增加了额外计算开销;极简的 Top-1 单专家激活即可实现兼顾效率与精度的最优表现。
- 权重归因剪枝保留核心轮廓:剪枝分析表明,按照 Grad-CAG 得分剪除 60% 最不重要的权重后,图像仍能较好保留主要目标物体的宏观结构;在 Pointing Game 目标定位指标上,Grad-CAG 取得 0.35,大幅击败基于幅值剪枝的 0.22 与随机剪枝的 0.15。
- 元学习消除流形割裂:UMAP 流形分析表明,未经元学习反向传播的 HMoE-WT 呈现破碎离散的类簇且泛化测试漂移显著,而 HMoE-MWT 形成了平滑连续的多类共享功能子电路流形。
亮点与洞察¶
- 拓扑先验与条件计算的无缝契合:提出针对 INR 权重组织特性的双阶段 MoE 结构,用层级均值路由捕捉全局频域特性,用跨层偏置投影指导局部 Token 专家分流,在保持与密集模型相同激活算力(\(k=1\))的同时实现了更强的表征解耦。
- 首个权重空间专用的可解释性分析套件:提出了 Grad-CAG 梯度敏感度度量,摆脱了以往仅能在像素空间做解释的局限,不仅能够通过剪枝量化检验子电路功能完整性,还能灵敏感知出 Imagenette 中“丁鱥鱼(tench)与渔网频繁共现”等数据集虚假关联偏置。
- 大分辨率可扩展性验证:在全量 ImageNet-1K 上首次将纯权重空间分类器的准确率推高至 26.73%,证明了条件计算在极端高维复杂函数拟合任务中的伸缩潜力。
局限与展望¶
- 单样本拟合的计算前置瓶颈:尽管元学习的 \(k\) 步内循环大幅平摊了拟合开销,但为每个新增图像在推理前必须独立优化一份 INR 权重的两阶段机制,仍然阻碍了在工业实时场景中的即时部署。
- 模型覆盖面局限于图像 SIREN:当前验证主要基于以正弦激活函数为核心的 2D 图像隐式神经表示,尚未延伸至具有更复杂几何结构的 3D NeRF、3D 高斯泼溅(3DGS)或周期多波形音频领域。
- 权重归因方法仍处于初级阶段:Grad-CAG 仅依赖一阶局域梯度与绝对值乘积,缺乏高阶交互效应(如积分梯度或二阶 Hessian 交互)建模,对深度网络参数冗余环境下的全局归因精度仍有提升空间。
相关工作与启发¶
- vs ScaleGMN / DWS-Net / NFN: 早期工作着重在分类器内部设计严格的神经元排列与缩放等变层(Equivariant Layers),但此类架构通常计算复杂度呈超线性爆发,难以扩展到高分辨率大网络。本文跳出单纯的等变数学约束,采用元学习规范化权重分布,并引入高效的 MoE 架构,在大规模数据集上展现了压倒性的伸缩性。
- vs Meta Weight Transformer (MWT): MWT 首次打通了元学习与权重 Transformer 的联合优化,但在分类器骨干中采用了平铺直叙的密集前馈网络。本文引入分层与局部两级 MoE 路由机制,在完全不增加单 Token 推理 FLOPs 的前提下,不仅 CIFAR-10 提升超 6%,而且为权重空间补充了不可或缺的归因解释工具。
评分¶
- 新颖性: ⭐⭐⭐⭐ [首次将分层 MoE 条件计算与一阶梯度归因引入 INR 权重分类,框架设计贴合网络内在拓扑]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖从 MNIST 到 ImageNet-1K 的多尺度基准,且包含详尽的参数等量对比、消融研究与剪枝定位验证]
- 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密连贯,图表呈现规范详实,数学推导与实验细节交待清晰]
- 价值: ⭐⭐⭐⭐ [为隐式表示学习、权重空间模式识别与神经网络白盒化解释提供了极具启发性的技术路线]