🤖 machine learning

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

该论文提出专家选择（EC）路由是扩散语言模型（DLM）中优于令牌选择（TC）路由的混合专家架构方案，通过引入基于去噪步数的自适应专家容量分配机制，实现了负载平衡优化、收敛加速及性能提升，并证明了现有 TC 模型仅需替换路由模块即可无缝迁移至 EC 范式。

原作者： Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia, Yang Zhou, Guanhua Chen, Ming Liu

发布于 2026-04-03

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia, Yang Zhou, Guanhua Chen, Ming Liu

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文主要讲的是如何给一种新型的人工智能语言模型（叫做“扩散语言模型”）装上更聪明的“大脑调度系统”，让它学得更快、跑得更快、还更省电。

为了让你轻松理解，我们可以把训练这个 AI 模型想象成指挥一个超级庞大的交响乐团排练一首复杂的曲子。

1. 背景：旧模式 vs. 新模式

旧模式（Token-Choice，令牌选择）：
想象一下，乐团里有几百种不同的乐器（专家），还有成千上万个乐谱片段（Token/单词）。
在旧模式下，每个乐谱片段自己决定找哪个乐器来演奏。

问题： 大家都喜欢找小提琴手，结果小提琴手累得半死，而大提琴手却在那儿发呆。
后果： 指挥（系统）必须停下来等最忙的小提琴手，导致整个乐团排练效率极低，而且为了平衡大家的工作量，还得额外搞一些复杂的“纪律惩罚”（辅助损失函数），既费时间又容易出错。

新模式（Expert-Choice，专家选择）：
这篇论文提出，不如反过来：让乐器（专家）自己来挑乐谱片段。

做法： 每个乐器手规定：“我只接 10 个片段，谁先来我挑谁，满了就停。”
好处： 这样每个乐器手的工作量是绝对平均的。没有人会累死，也没有人会闲着。整个乐团可以齐步走，效率直接翻倍。

2. 核心发现：什么时候该“用力”？

既然每个乐器手的工作量固定了，那能不能更聪明一点？比如，在曲子最难的部分多派点人手，简单的部分少派点？

作者发现，在 AI 学习的过程中，有一个叫“去噪步数”（Timestep）的概念。

高噪点（Mask Ratio 高）： 就像乐谱全是乱码，几乎看不清字。这时候派再多专家，大家也是瞎猜，效率很低。
低噪点（Mask Ratio 低）： 就像乐谱已经写好了 90%，只剩下几个词没填。这时候专家只要稍微动脑筋，就能精准补全，学习效率极高。

作者的策略（线性反向调度）：
作者设计了一个聪明的调度员：

当乐谱还很乱（高噪点）时，只派少量专家去“扫盲”。
当乐谱快写好了（低噪点）时，派大量专家去“精修”。

比喻： 这就像你写论文。

刚开始头脑一片空白（高噪点），你花 1 小时可能只能写出 1 句话。
快写完时，思路很清晰（低噪点），你花 1 小时能写出 10 页。
这篇论文建议：把大部分精力（算力）都花在快写完的那段黄金时间，而不是死磕开头。结果发现，这样不仅写得快，质量还更高。

3. 三大成果

速度快得惊人：
用新系统（专家选择）训练模型，比旧系统快 2 倍。就像把原本需要 20 小时的排练，缩短到了 10 小时。
不浪费资源：
旧系统经常因为“木桶效应”（最慢的那个专家拖慢所有人）而卡顿。新系统让所有人步调一致，GPU 显卡的内存使用像切蛋糕一样均匀，没有浪费。
旧模型也能“换芯”：
最棒的是，对于那些已经训练好的、用旧系统的模型，不需要重新训练，只要把“调度员”（路由器）换掉，就能立刻获得速度提升和效果优化。这就像给老式汽车换了一个更聪明的变速箱，车还是那辆车，但跑起来更顺了。

总结

这篇论文的核心思想就是：别让 AI 像无头苍蝇一样乱撞（旧模式），要让它像精明的指挥官一样，在关键时刻（低噪点）集中火力，平时（高噪点）保持节奏。

通过这种“专家选择”和“动态分配”的策略，未来的 AI 模型不仅能变得更聪明，还能在训练和使用时更省钱、更快速。

这是一篇关于扩散语言模型（Diffusion Language Models, DLMs）中混合专家（MoE）架构路由机制的预印本论文。论文提出并验证了专家选择路由（Expert-Choice, EC）优于传统的令牌选择路由（Token-Choice, TC），并进一步提出了基于时间步的自适应专家容量调度策略。

以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

背景：扩散语言模型（DLMs）通过迭代去噪实现并行非自回归文本生成，已成为自回归模型的重要替代方案。为了扩展模型规模，DLMs 也开始采用混合专家（MoE）架构。
现有问题：
- 负载不平衡：现有的 DLM MoE 模型直接沿用了自回归系统中的**令牌选择（Token-Choice, TC）**路由机制（即每个令牌独立选择其偏好的专家）。由于令牌是独立选择的，导致部分专家过载而其他专家空闲，造成严重的负载不平衡。
- 计算僵化：TC 路由需要引入辅助的负载平衡损失（Auxiliary Load-Balancing Loss），这不仅消耗额外计算资源，还可能干扰主语言建模目标的梯度。即使如此，负载平衡依然脆弱，无法提供硬性保证。
- 未利用 DLM 特性：DLM 在每一步去噪过程中是非因果的（Non-causal），且同时处理所有令牌，这与 TC 路由所需的“令牌视角”不匹配，反而更适合“专家视角”。
- 缺乏自适应计算：DLM 的训练和推理涉及不同的去噪时间步（Timesteps），每个时间步的掩码比率（Masking Ratio）不同，任务难度也不同。现有的 TC 路由无法根据时间步动态调整计算量。

2. 方法论 (Methodology)

2.1 专家选择路由 (Expert-Choice Routing, EC)

核心机制：将路由逻辑反转。在 EC 中，每个专家选择固定数量（ $c$ ）的令牌，而不是令牌选择专家。
优势：
- 确定性负载平衡：由于每个专家处理的令牌数量是固定的，天然消除了负载不平衡，无需辅助损失函数。
- 全局视图：DLM 在每一步去噪时都能访问所有令牌，完美契合 EC 所需的“专家侧选择”机制。
- 吞吐量提升：消除了因负载不均导致的 GPU 等待（Straggler effect），显著提高了训练吞吐量。

2.2 时间步自适应专家容量调度 (Timestep-Dependent Expert Capacity)

动机：DLM 在不同去噪时间步面临的掩码比率（ $r$ ）不同。作者假设不同时间步的学习效率不同，因此计算资源分配不应是固定的。
策略：将专家容量 $k$ $k$ 定义为掩码比率 $r$ $r$ 的函数 $k(r)$ $k (r)$ 。
- 提出了多种调度策略（如线性反向、余弦反向、高斯分布等），旨在将更多计算资源分配给特定的掩码比率区间。
- 线性反向调度（Linear-Reverse）：随着去噪过程进行（掩码比率 $r$ 降低），增加专家容量。即在低掩码比率（接近完成）阶段分配更多专家。

2.3 机制分析

通过测量不同掩码比率区间的收敛速率（Convergence Rate, $\eta_r$ ），发现低掩码比率（Low-mask-ratio）阶段的令牌学习效率比高掩码比率阶段高出一个数量级。
结论：在低掩码比率阶段，大部分令牌已可见，模型拥有丰富上下文，此时增加专家容量能带来最大的边际收益；而在高掩码比率阶段，上下文匮乏，增加专家带来的收益递减。

2.4 现有模型改造 (Retrofitting)

证明了现有的预训练 TC-DLM 模型可以通过**仅替换路由模块（Router）**的方式，无缝转换为 EC 架构，而无需重新训练专家权重或嵌入层。

3. 关键贡献 (Key Contributions)

确立 EC 为 DLM 的优选路由范式：通过从头训练实验证明，EC 路由在负载平衡、吞吐量（比 TC 快 1.5-2.1 倍）和收敛速度（比 TC 快 2.0 倍）上均全面优于 TC 路由。
提出时间步自适应容量调度：首次引入基于时间步的动态专家容量分配。实验表明，将更多计算资源分配给低掩码比率阶段（Linear-Reverse 策略）在匹配 FLOPs 的情况下，能显著降低验证困惑度（Perplexity）。
揭示学习效率机制：通过机制分析发现，低掩码比率步骤的学习效率最高，解释了为何动态调度有效。
验证模型改造可行性：展示了将预训练的 TC-DLM 改造为 EC-DLM 的简单性和有效性，改造后模型在微调（SFT）阶段收敛更快，推理速度提升 1.3-1.5 倍，且精度提升。

4. 实验结果 (Results)

训练效率：
- 在相同架构和数据下，EC 达到损失 3.75 仅需 10.6 小时，而 TC 需要 20.7 小时（快 2 倍）。
- EC 的吞吐量达到 52.1 TFLOP/s/GPU，显著高于 TC 的各种变体（24.9 - 35.4 TFLOP/s/GPU）。
- GPU 显存使用在 EC 下完全均匀（标准差 0.0 GB），而 TC 下波动巨大（标准差 3.6 GB）。
调度策略对比：
- 在 OpenWebText 和 Nemotron-CC 数据集上的实验显示，Linear-Reverse（低掩码比率高容量）策略表现最佳，优于静态容量和偏向高掩码比率的策略。
- 在 8B 参数规模的预训练中，动态 EC 在验证困惑度、MMLU 和 ARC-Challenge 等下游任务上均优于静态 EC。
微调与推理：
- 将预训练的 LLaDA-MoE（TC 架构）改造为 EC 后，在 HumanEval、GSM8K、MedQA 等任务上，收敛速度显著加快。
- 推理加速：由于消除了负载不平衡，EC 的解码速度比 TC 快 1.3-1.5 倍。
- 精度提升：动态 EC 在多个基准测试中达到了最高的最终精度（例如 HumanEval Pass@1 从 53.9% 提升至 58.6%）。

5. 意义与影响 (Significance)

范式转变：该工作证明了在扩散语言模型中，计算资源不应被视为固定的架构常数，而应作为一种自适应策略（Adaptive Policy）。
解决核心瓶颈：通过 EC 路由彻底解决了 MoE 模型中的负载不平衡问题，无需辅助损失，简化了训练流程并提升了硬件利用率。
可落地性：提出的“仅替换路由”的改造方案，使得现有的大规模预训练 TC-DLM 模型可以低成本地升级为更高效的 EC 架构，具有极高的工程应用价值。
理论洞察：揭示了 DLM 在不同去噪阶段的学习动态差异，为未来设计更精细的扩散模型调度策略提供了理论依据。

总结：这篇论文通过引入专家选择路由（EC）和基于时间步的自适应容量调度，显著提升了扩散语言模型的训练效率、推理速度和最终性能，为 DLM 的规模化应用提供了新的技术路径。代码已开源。