动态稀疏专家混合的分布一致推理

专家混合 (MoE) 架构已成为扩展模型容量的强大范例，同时在大型基础模型中保持高效推理。然而，大多数 MoE 模型使用固定的 top-k$ 专家选择策略，为每个词元分配相同的专家预算，即使较少的专家可能就足够了。推理时动态 top-$k$ 路由可以减少计算量而无需重新训练，但现有方法经常忽略由于偏离训练时路由配置而导致的分布变化。我们表明，减少激活专家的数量会持续增加 SMoE 输出的 RMS 规模和方差，从而引起表示不匹配，除了专家容量损失之外，还会导致下游性能下降。为了解决这个可校正组件，我们提出了逐层分布对齐（LDA），这是一种轻量级推理时间校正，它使用逐层校准统计数据将简化的路由表示与默认配置对齐。在多个 SMoE LLM、基准测试和路由策略中，LDA 恢复了减少路由下分布偏移引起的大部分性能损失，同时以可忽略的开销保持稀疏推理效率。