[Paper] UniPool：面向 Mixture-of-Experts 的全局共享专家池

Source: arXiv - 2605.06665v1

Overview

本文介绍了 UniPool，这是一种对 Mixture‑of‑Experts（MoE）模型的新改进，它用 单个全局共享的专家池 取代了传统的“每个 transformer 层对应一组专家”的规则。通过让每一层都从同一个池中挑选专家，作者展示了在减少专家参数总量的同时仍能提升模型质量——这对可扩展性和效率都是一次双赢。

关键贡献

• 全局专家池: 用一个共享池取代每层的专家集合，由每层的独立路由器访问。
• 平衡利用率损失: 一种新颖的辅助损失，鼓励整个池子被均匀使用，防止少数专家占用过多流量。
• NormRouter: 一种尺度稳定的路由机制，能够可靠地与共享池配合，保持稀疏性和训练稳定性。
• 实证收益: 在五种 LLaMA‑style 模型规模（182 M–978 M 参数）上，使用 30 B 令牌进行训练，UniPool 始终降低验证损失（相对最高‑0.0386）和困惑度，相较于普通 MoE 基线。
• 参数效率: 证明池的规模可以视为深度缩放的超参数；缩小池子（仅使用 41.6 %‑66.7 % 的专家预算）的变体能够匹配或超越传统层级 MoE。
• 可组合性: 表明 UniPool 的优势可以与更细粒度的专家分解技术叠加。

方法论

1. 共享专家池设计

   • 作者没有为每个 Transformer 层分配独立的专家集合，而是创建了一个包含 N 个专家的单一池。
   • 每层仍保留自己的路由器，但路由器现在从 全局池 中选择专家，而不是从私有子集选择。

2. 路由机制（NormRouter）

   • 实现了归一化评分函数，能够在保持对隐藏状态幅度变化的鲁棒性的同时，产生稀疏的 top‑k 选择。
   • 这避免了“尺度漂移”问题，否则该问题可能导致某些层在池中占主导。

3. 池级辅助损失

   • 计算整个池在一个批次中每个专家的负载（即处理的 token 数量）。
   • 添加一个正则化项，惩罚偏离均匀负载分布的情况，从而鼓励专家使用的平衡。

4. 训练设置

   • 模型遵循 LLaMA 架构，参数规模从 1.82 亿到 9.78 亿不等。
   • 在 Pile 数据集上使用 300 亿 token 进行训练，采用与原始 MoE 基线相同的 token 预算和训练计划，以确保公平比较。

5. 评估

   • 在保留集上测量验证损失和困惑度。
   • 额外的消融实验测试了不同的池大小以及与其他专家划分策略的兼容性。

结果与发现

• 均衡利用率： 辅助损失成功分散流量；相较于 vanilla MoE，专家负载的标准差下降约 45 %。
• 参数节省： 将专家池削减至原始规模约 50 % 的 UniPool 变体仍能实现 相等或更好 的损失，证明专家容量不必随深度线性增长。
• 可组合性： 与 token‑level 专家门控或层次专家划分等技术结合时，UniPool 还能额外降低 0.01–0.02 的损失，显示出叠加效益。

Practical Implications

• Reduced Memory Footprint: 共享专家会减少权重矩阵的总数量，从而降低 GPU 内存需求——这在受限硬件上进行大规模 LLM 训练时尤为宝贵。
• Faster Inference: 更小的专家池意味着更少的内存访问和更好的缓存局部性，进而降低 MoE‑enabled 服务（如代码补全、聊天助手）的延迟。
• Simplified Scaling Strategy: 工程师可以将 pool size 视为单一超参数，而不必在每层之间分别调节专家数量，使模型扩展更加可预测。
• Cost‑Effective Training: 参数同步量减少后，分布式训练成本更低，带宽需求也更友好。
• Flexibility for Multi‑Task Settings: 全局专家池可以自然地在不同任务或领域之间复用，形成共享的“专家知识库”，无需在每层重复复制。

限制与未来工作

• Routing Overhead: 虽然 NormRouter 稳定了训练，但全局负载均衡的额外计算相较于最简单的 top‑k 路由器会带来适度的开销。
• Scalability Beyond 1 B Parameters: 实验止于小于 1 B 参数的模型；尚不清楚 UniPool 在多十亿参数的大型语言模型中表现如何，因为那时通信成本占主导。
• Specialized Expert Needs: 某些层（例如早期嵌入层与深层推理层）可能受益于 层‑特定 的专长；完全共享的池可能会稀释这种专门化。
• Future Directions: 作者建议探索层次化池（例如每阶段共享池）、训练期间动态池大小调整，以及将 UniPool 与检索增强或基于适配器的微调流水线结合。

作者

• Minbin Huang
• Han Shi
• Chuanyang Zheng
• Yimeng Wu
• Guoxuan Chen
• Xintong Yu
• Yichun Yin
• Hong Cheng

论文信息

• arXiv ID: 2605.06665v1
• 分类: cs.LG, cs.AI
• 发表时间: 2026年5月7日
• PDF: 下载 PDF