arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2603.11535cs.AIcs.CL
收藏

专家阈值路由用于自回归语言模型的动态计算分配与负载均衡

Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing

Hanchi Sun, Yixin Liu, Yonghui Wu, Lichao Sun

AI总结:

本文提出专家阈值路由方法,通过动态计算分配和负载平衡机制提升自回归语言模型性能,实验显示其在预训练任务中表现更优。

摘要:

Token-choice Mixture-of-Experts (TC-MoE) 为每个 token 分配固定数量的专家,限制了动态计算分配并需要辅助损失来维持负载平衡。我们提出了专家阈值 (ET) 路由,其中每个专家维护一个从全局 token 分布估计出的指数移动平均 (EMA) 阈值。在训练和推理过程中,每个 token 如果其分数超过专家的阈值,则独立路由到该专家,从而实现动态计算分配,同时在不使用辅助损失的情况下实现负载平衡。这种完全因果机制消除了对批量中其他 token 的依赖,使其非常适合自回归语言模型。在预训练实验中,针对 FineWeb-Edu 数据集扩展到 2.4B 参数时,ET 的交叉熵损失比 TC-MoE 低 0.067,相当于用 1.6 倍更少的 token 达到相同性能。

英文摘要:

Token-choice Mixture-of-Experts (TC-MoE) routes each token to a fixed number of experts, limiting dynamic computation allocation and requiring auxiliary losses to maintain load balance. We propose Expert Threshold (ET) routing, where each expert maintains an exponential moving average (EMA) threshold estimated from the global token distribution. At both training and inference, each token is independently routed to an expert if its score exceeds the expert's threshold, enabling dynamic computation allocation while achieving load balance without auxiliary losses. This fully causal mechanism eliminates dependence on other tokens in the batch, making it well-suited for autoregressive language modeling. In pretraining experiments scaling to 2.4B parameters on FineWeb-Edu, ET achieves 0.067 lower cross-entropy loss than TC-MoE, equivalent to reaching the same performance with 1.6$\times$ fewer tokens.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑