专家选择路由支持扩散语言模型中的自适应计算

扩散语言模型 (DLM) 支持并行、非自回归文本生成，但现有的 DLM 专家混合 (MoE) 模型继承了自回归系统的标记选择 (TC) 路由，导致负载不平衡和严格的计算分配。我们表明，专家选择 (EC) 路由更适合 DLM：它通过设计提供确定性负载平衡，比 TC 产生更高的吞吐量和更快的收敛。基于 EC 能力是外部可控的特性，我们引入了时间步长相关的专家能力，它根据去噪步骤改变专家分配。我们发现，为低掩码比步骤分配更多容量可以在匹配的 FLOP 下始终实现最佳性能，并提供机械解释：低掩码比上下文中的词元表现出更高数量级的学习效率，因此将计算集中在这些步骤上会产生最大的边际回报。最后，我们表明，可以通过仅替换路由器来将现有的预训练 TC DLM 改造为 EC，从而实现更快的收敛并提高不同下游任务的准确性。总之，这些结果将 EC 路由确立为 DLM MoE 模型的高级范例，并证明 DLM 中的计算可以被视为自适应策略而不是固定的架构常量。代码可在 https://github.com/zhangshuibai/EC-DLM 获取。