用专家选择路由提升扩散语言模型的计算效率和平衡性。
Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models
- 采用专家选择路由,实现确定性负载均衡。
- 低掩码率步骤分配更多算力,性能最优且节省计算。
- 可替换现有模型路由器,适配多种下游任务。
扩散语言模型(DLM)支持并行非自回归文本生成,但现有混合专家(MoE)模型沿用自回归系统的词元选择(TC)路由,导致负载不均与计算分配僵化。本文表明专家选择(EC)路由更适用于DLM:其设计上保证确定性负载均衡,带来更高吞吐量和更快收敛。基于EC容量可外部调控的特性,我们提出随时间步变化的专家容量机制,根据去噪步骤动态调整专家分配。实验发现,在相同浮点运算量(FLOPs)下,将更多容量分配给低掩码率步骤始终表现最佳,其机制解释为:低掩码率上下文中的词元学习效率高出一个数量级,因此在此阶段集中计算获得最大边际收益。最后,我们证明现有预训练的TC-DLM可通过仅替换路由器转化为EC-DLM,实现更快收敛与更优准确率,覆盖多种下游任务。结果确立了EC路由作为DLM MoE模型的更优范式,并揭示扩散语言模型的计算可视为一种自适应策略而非固定架构。代码已开源:https://github.com/zhangshuibai/EC-DLM。
原文摘要 · Abstract (English)
Diffusion language models (DLMs) enable parallel, non-autoregressive text generation, yet existing DLM mixture-of-experts (MoE) models inherit token-choice (TC) routing from autoregressive systems, leading to load imbalance and rigid computation allocation. We show that expert-choice (EC) routing is a better fit for DLMs: it provides deterministic load balancing by design, yielding higher throughput and faster convergence than TC. Building on the property that EC capacity is externally controllable, we introduce timestep-dependent expert capacity, which varies expert allocation according to the denoising step. We find that allocating more capacity to low-mask-ratio steps consistently achieves the best performance under matched FLOPs, and provide a mechanistic explanation: tokens in low-mask-ratio contexts exhibit an order-of-magnitude higher learning efficiency, so concentrating compute on these steps yields the largest marginal return. Finally, we show that existing pretrained TC DLMs can be retrofitted to EC by replacing only the router, achieving faster convergence and improved accuracy across diverse downstream tasks. Together, these results establish EC routing as a superior paradigm for DLM MoE models and demonstrate that computation in DLMs can be treated as an adaptive policy rather than a fixed architectural constant. Code is available at https://github.com/zhangshuibai/EC-DLM.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。