UniPool：面向混合专家模型的全局共享专家池 UniPool: A Globally Shared Expert Pool for Mixture-of-Experts

UniPool 的训练流程可以分为四步。第一步，模型实例化：在 Megatron-LM 中只实例化一份专家模块 $\mathcal{E}=\{e_1,\dots,e_M\}$，所有 MoE 层共享同一组权重指针，但每层保留独立的路由器 $r_l$，路由器的参数与 Transformer 注意力都按层独立。这是与「逐层 MoE」唯一的结构差异。第二步，路由打分：对每个 token 的 hidden state $h$，路由器计算 logits $z=Wh$，再按 NormRouter 公式 $s_i=\sigma\cdot c\cdot\max(0,z_i/(\|z\|_2+\epsilon))$ 得到每个共享专家的分数，对分数取 Top-k 激活。第三步，池级均衡损失计算：每个 micro-batch 内收集所有 L 个共享层在共享池上的 token 分配直方图 $f_i^{(l)}$，按 $\bar f_i=\frac{1}{L}\sum_l f_i^{(l)}$ 聚合为全局分配，加权路由概率 $\bar P_i=\frac{1}{L}\sum_l P_i^{(l)}$ 同理聚合，得到池级辅助损失 $\mathcal{L}_{\text{pool}}=\alpha_{\text{pool}}\cdot M\cdot\sum_{i=1}^{M}\bar f_i\bar P_i$；为避免跨层张量依赖，实际实现把全局统计延后一个 micro-batch 计算（细节见 Appendix G），主损失仍是各层独立贡献之和。第四步，主损失整合：训练总损失 $\mathcal{L}=\mathcal{L}_{\text{NTP}}+\mathcal{L}_{\text{pool}}$，不附加任何逐层辅助损失；优化器用 AdamW + cosine 学习率 + bf16 Megatron-LM；训练数据为 Pile，60K 步、batch size 512、序列长 1024，约 30B tokens。实验配置分主表（5 个规模，8E/top-1）与粒度扫描（182M 上 16E/top-2、32E/top-4）两组，并配套在 182M 上做组件消融（共享池 + 池级 aux + NormRouter 三件套逐一开/关），以及共享范围 G（1=全局，12=逐层）的中间扫描。

UniPool 的技术新颖性主要体现在三个层面。第一层是架构层面：MoE 设计领域长期被「逐层独立专家」这一默认假设支配，作者用一个全局共享池 + 独立路由器打破了这条假设，并且明确指出这种「按所有权粒度对齐均衡目标」是此前工作（包括 MoEUT、ALBERT 等跨层共享路线）所缺乏的针对性改动。第二层是方法层面：池级辅助损失的设计看似简单（把 $f_i^{(l)}$ 替换为 $\bar f_i$），但物理意义深远——它把「避免死专家」这件事从「每层局部」重新定义为「全局参数利用」，并通过 Figure 3 的实证可视化展示，如果仍用逐层损失，共享池会全局塌缩到极少数专家（Figure 3a），而 UniPool 则恢复均匀利用（Figure 3b）；把 NormRouter 引入是为了解决大池下 softmax 不稳定的问题，这一「池级损失 + NormRouter」的协同设计是本文最核心的方法学贡献。第三层是结论层面：通过 Figure 2(a) 的 reduced-pool 扫描证明专家参数可以随深度亚线性增长（182M 用 66.7% / 469M 用 50% / 830M 用 41.6% 即可反超逐层 MoE），并通过 830M vs 978M 对比展示「向深度 + 共享池分配参数」比「向宽度分配参数」更高效（830M 的 UniPool 1.6923 优于 978M 的 1.6999，尽管后者参数量更大）——这为 MoE 缩放法则打开了一个新维度：池大小可以作为一个独立的深度缩放旋钮。

从独立视角看，UniPool 仍存在几个值得改进的弱点。第一，对池大小的「甜区」缺乏系统搜索：Figure 2(a) 只在 4 个规模上扫了 3–4 个池大小点，没给出「最小可接受池大小」与深度的显式函数关系或拟合公式，工程上很难直接据此设定新模型的池大小——可以引入一个池大小预测模型或基于深度的 closed-form 启发式。第二，池级辅助损失在 micro-batch 延后聚合的实现在分布式训练（特别是 expert-parallel 数据并行）下需要跨 rank 同步 $\bar f_i$，当前论文未给出多机多卡下的通信开销与同步策略，建议补一份分布式 benchmark。第三，作者声称 UniPool 与细粒度 MoE 兼容（Figure 2b），但只在 182M 上验证；如果专家数从 32 扩到 128 或 256、池大小相应从 384 扩到数千，NormRouter 的 L2 归一化在极大池上是否仍能保持稀疏性（约半数得分为 0）需要验证，可能需要配合 expert-choice 路由或 top-k 截断做进一步协同设计。第四，路由随机化探针（Table 4）的 UniPool 列用的是 cardinality-matched top-8 随机，与生产模型的「整层随机」不可直接比较，作为反冗余证据不够干净；理想的做法是同时报告 full-pool uniform random 与 cardinality-matched random 两组数值，让读者看清「池冗余度下降」与「路由决策权重增加」两条因果链。第五，论文没有分析 UniPool 在推理时的 memory access pattern：所有 L 层共享同一份大专家权重，意味着每次前向要在更大的权重矩阵上做索引访问，可能破坏逐层 MoE 的 cache locality；建议补一份 inference latency / HBM bandwidth 实测。第六，单一 Pile 数据 + 单一任务（next-token prediction）下评估，没有覆盖 instruction tuning、RLHF、长上下文、tool use 等场景，论文生态适配性证据不足。

作者在文中提出了若干明确方向：第一，把池大小正式建模为深度缩放超参数，并拟合「层数 L ↔ 最优池大小 M」的经验缩放律，挑战 MoE 「专家数 = 层数 × 每层专家」的默认设定；第二，把 UniPool 扩展到更大规模（7B+）与多模态场景，验证池级损失在视觉专家 / 音频专家共存下是否仍稳定；第三，进一步把专家共享思想泛化到「细粒度 + 跨层共享」的复合 MoE。我的延伸建议包括：第一，把 NormRouter 的可学习标量 $\sigma$ 升级为层相关甚至 head 相关的向量，让每层路由器能更细粒度地控制路由锐度；第二，结合 expert-choice 路由，让每个共享专家主动选 token 而非被动被打分，可能解决 Figure 3(a) 那种「专家选择权不对等」导致的池塌缩；第三，探索「混合所有权」——例如底层 6 层共享一个池、中间 6 层共享另一个池、顶层 6 层逐层私有——把 Section 6.2 中「unique expert weights touched」从 94.1%/89.5%/82.7% 这条单调下降曲线变成可学习的曲线；第四，把池级辅助损失与 router z-loss（ST-MoE）联合使用，进一步抑制 logit 漂移；第五，理论化分析：在何种「专家相似度分布」下共享池严格不弱于逐层私有，给出可解释的池大小下界。

论文的可复现性较好，体现在以下几点：第一，代码已开源在 https://github.com/Centaurus-Alpha/UniPool，主干实现基于 Megatron-LM（与 ST-MoE / GShard 等主流 MoE 实现同源），专家池模块被实例化一次并在 MoE 层间共享，改动面较局部；第二，主表所有模型（5 个规模）使用完全公开的 Pile 数据集 [15]，训练超参（AdamW、cosine LR、bf16、batch size 512、序列长 1024、60K 步 ≈ 30B tokens）在 Appendix D 完整给出；第三，模型架构遵循 LLaMA 规范，5 个规模的层数 / 隐藏维 / 注意力头 / FFN 维度在 Appendix Table 6 完整列出；第四，主要路由配置（8E/top-1）作为默认协议，附加 16E/top-2、32E/top-4 粒度扫描实验在 182M 上公开。潜在的复现难点包括：第一，182M 跑了 3 个随机种子取平均，更大规模仅 1 次运行，数值波动范围未给出置信区间；第二，池级辅助损失需要 micro-batch 延后聚合（Appendix G），实现细节对最终数值有影响，必须严格按论文伪代码实现；第三，作者未公开 NormRouter 中常数 $c$ 的具体蒙特卡洛采样数与随机种子，附录 H 只给了数学期望与采样流程描述；第四，完整训练一张 830M 模型所需的 GPU 小时数与硬件配置未明示，第三方复现需要相当算力预算（保守估计至少需要数十张 A100 / H100 连续运行数天），不适合算力有限的研究组全规模复现，但 182M + 共享池 + 池级 aux + NormRouter 的最小验证集可在单卡数小时内跑完。