专家选择路由为扩散语言模型带来自适应计算能力 Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models
EC路由+时间步容量调度让DLM MoE训练提速2倍、推理提速1.5倍
前置知识
扩散语言模型(DLM)
DLM 通过迭代去噪生成文本：从全掩码序列出发，按掩码调度 $\gamma(t)$ 逐步预测并解掩码位置得到完整序列。每步用双向注意力同时处理所有 $L$ 个位置，与自回归逐 token 生成有本质区别。
DLM 的非因果、整序列并行去噪特性是本文一切论断的根基。理解 DLM 是理解为什么 EC 路由（需要全局视图）在 DLM 中自然适配的前提。
混合专家(MoE)与稀疏路由
MoE 把单个 FFN 替换为 $E$ 个并行专家 FFN，由路由器对 $N$ 个 token 输出门控分数矩阵。稀疏路由只激活部分 token-expert 对，TC 让 token 选 expert，EC 让 expert 选 token。
理解 TC 与 EC 在'谁选谁'上的对称反转是本文核心。TC 的负载不平衡是 DLM MoE 训练的瓶颈，本文通过反转选择方向同时获得吞吐提升和新的调度自由度。
Token-Choice(TC)路由
TC 让每个 token 独立选 top-$k$ 个 expert：$E_i=\mathrm{TopK}_j(S_{i,j},k)$。由于 token 独立决策，专家负载不可控，需要容量因子与辅助负载均衡损失双重约束，存在 straggler 效应。
TC 是基线。TC 的负载不均与容量上限两个缺陷通过 per-GPU memory std 3.6 GB 等实验被量化，是评估 EC 收益的对照基础。
Expert-Choice(EC)路由
EC 反转选择方向：每个 expert 独立选 top-$c$ 个 token，$T_j=\mathrm{TopC}_i(S_{i,j},c)$。设 $c=kN/E$ 时与 TC top-$k$ 计算量匹配，但负载均衡由构造保证，无需辅助损失，$c$ 本身成为可调设计变量。
EC 是本文的核心替代方案。理解其'专家选 token'的反转逻辑以及'负载由构造决定'的性质，是理解 EC 既消除 straggler 又开放时间步容量调度的关键。
时间步自适应计算
在扩散模型中，不同去噪步对应不同掩码率 $r$，每步面临质性不同的子任务。自适应计算指根据当前步特性动态分配计算资源（如专家容量），形式化为 $k(r)=k_{\min}+(k_{\max}-k_{\min})\cdot s(r)$。
时间步自适应是本文第二个核心创新。理解为什么 DLM 的不同掩码率步需要不同计算量，以及为什么低掩码率步学习效率高一个数量级，是把握'把计算视为策略而非常量'的关键。
研究动机
扩散语言模型(DLM)如 LLaDA、MDLM 已发展为替代自回归的前沿范式，并通过 MoE 扩展至百亿参数规模。然而现有 DLM MoE 直接沿用自回归系统的 token-choice (TC) 路由，带来三个问题。**结构性负载不均衡**：TC 让每个 token 独立选 top-$k$ expert，per-expert 负载不可控。在 LLaDA-2.0-mini (16B) 的 8 卡 H100 推理中，per-GPU 显存标准差 3.6 GB，单卡占 70.3 GB 而其他卡仅 58–64 GB，产生 straggler 效应。**辅助损失副作用**：TC 需额外负载均衡损失 $\mathcal{L}_{\mathrm{LB}}$，引入与主 LM 目标竞争的梯度信号；消融实验显示 dropless TC 配合 loss-free bias 也无法弥合 1.7 倍速度差距，说明负载不均是结构性问题。**计算调度僵化**：TC 下 per-expert 负载是 token 独立选择的涌现结果，无法被外部控制；DLM 中'不同去噪步面临不同任务'的结构特性无法被利用。
本文的目标是本文的具体目标是双重的。第一个目标是系统验证 expert-choice (EC) 路由作为 DLM MoE 替代方案的优越性——通过严格控制架构、数据、超参数相同的对照实验，量化 EC 在吞吐、收敛速度、负载均衡三个维度上相对 TC 的提升幅度。第二个目标是探索并形式化'时间步自适应计算'这一新维度：在 EC 路由下，expert 容量 $c$ 成为可外部调节的设计变量，作者据此提出容量调度函数 $k(r)=k_{\min}+(k_{\max}-k_{\min})\cdot s(r)$，系统比较多种调度策略（线性、cosine、Gaussian 及其反向），并在 FLOPs 匹配的前提下找出最优调度。最终，作者还希望证明：现有部署的 TC 预训练 DLM（如 LLaDA-MoE）可以通过仅替换路由器的方式无缝迁移到 EC，从而把研究成果直接惠及现有模型。
与已有工作不同的是，本文独特切入角度有三。**路由机制与扩散结构特性的耦合**：之前工作要么沿用 AR 范式 TC 路由忽视 DLM 非因果特性，要么像 EC-DIT 仅在视觉连续扩散上验证 EC 路由。本文首个把 EC 路由系统性引入离散掩码扩散语言模型，利用 DLM 非因果整序列并行性论证 EC 在 DLM 中比 AR 更'自然'。**时间步作为新自适应维度**：之前 MoE 自适应计算主要沿深度（early-exit、MoD），或沿视觉扩散 timestep 做隐式/学习式容量分配。本文首次以显式、hand-designed、解析可比较的调度函数系统研究 DLM 中专家容量随掩码率变化，并通过收敛率 $\eta_r=-\frac{d\ln L_r}{dt}$ 给出机理——低掩码率步收敛率高出高掩码率步一个数量级。**EC 改造从训练扩展到现有 TC 模型 retrofit**：仅替换路由器即可让 LLaDA-MoE 在 GSM8K、HumanEval、MedQA 上同时获得更快收敛、更高准确率与 1.3–1.5 倍推理加速，为已部署模型的迁移提供了低成本路径。
核心方法
本文方法由两条相互独立又互补的技术路线组成。**第一条：路由机制替换**。将 DLM MoE 中标准的 token-choice (TC) 路由器替换为 expert-choice (EC) 路由器——把每个 token 输出 $k$ 个专家得分改为每个 expert 在所有 token 的得分中选 top-$c$。直觉上相当于'商品挑货架'改为'货架挑商品'，每 expert 处理 token 数严格相等，消除 straggler。**第二条：时间步自适应容量调度**。在 EC 路由基础上，把恒定 $c$ 升级为关于掩码率 $r$ 的函数 $k(r)=k_{\min}+(k_{\max}-k_{\min})\cdot s(r)$，系统比较多种 $s(r)$。所有调度在 $r\sim\mathrm{Uniform}(0,1)$ 下 $\mathbb{E}[k(r)]=\bar{k}$，保证与静态基线 FLOPs 严格匹配。第三条路由器 retrofit：保留预训练 TC DLM 的 expert 权重与 embeddings，仅替换路由器做下游 SFT。
本文的核心创新点与已有方法的本质区别在于'把 DLM 的时间步结构转换为 MoE 的可调计算自由度'。传统观点认为 MoE 路由只是'选择机制'，TC 路由已经广泛使用，需要辅助损失约束负载。但本文揭示了两层新视角：第一层是 EC 路由的'结构性优势'——通过反转选择方向，负载均衡由构造保证（每 expert 处理 $c$ 个 token 严格相等），消除 straggler，per-GPU memory std 由 3.6 GB 降至 0.0 GB，吞吐从 24.9 提升到 52.1 TFLOPs/s/GPU；第二层是 EC 路由独有的'可调度性'——容量 $c$ 作为外部设计变量打开了时间步自适应的大门，这是 TC 路由下根本无法实现的（TC 的负载是涌现的）。本文还通过机理性指标 $\eta_r=-\frac{d\ln L_r}{dt}$ 量化发现低掩码率步学习效率高出 7 倍以上，从而解释了为什么线性反向调度（更多容量给低 $r$）效果最好。
方法步骤详情
本文方法按两条主线执行。**主线 A：路由机制替换与时间步容量调度。** 步骤 1：构造对照架构——16 层 Transformer、hidden 2048、64 fine-grained expert 的 MoE DLM，TC 与 EC 共享所有超参数（AdamW, WSD, lr=$2\times10^{-4}$）。步骤 2：定义 EC 路由——每个 expert 选 top-$c$ 个 token：$T_j=\mathrm{TopC}_i(S_{i,j},c)$。步骤 3：定义时间步容量调度函数 $k(r)=\mathrm{clamp}(k_{\min}+(k_{\max}-k_{\min})\cdot s(r))$，所有调度器匹配 FLOPs。步骤 4：训练 OpenWebText 与 Nemotron-CC 上的 8B-A1B 模型。**主线 B：TC 模型 EC retrofit**——冻结 LLaDA-MoE 所有 expert 权重与 embeddings，仅替换路由器；在 GSM8K-AUG-NL、OpenCodeInstruct、MedQA 上做 SFT。
技术新颖性
本文的技术新颖性体现在四个层面。第一，**问题定义上的新颖**：首次把 DLM 的'去噪步结构差异'与 MoE 的'专家容量'显式耦合，提出时间步自适应容量调度这一新问题，区别于已有工作（DiffMoE 视觉扩散的 learned predictor、AdaDiff/DyDiT 的 width/depth 调节）。第二，**路由机制层面的迁移创新**：在 DLM 中首次系统证明 EC 优于 TC，并通过控制辅助损失（no LB/loss-free bias）与容量因子（1.0/1.25/1.5）的消融，证明 EC 的优势源于结构性负载均衡而非辅助损失效应。第三，**机理性发现**：通过在掩码率四分位 bin 上计算对数损失收敛率 $\eta_r$，揭示低掩码率步学习效率高出高掩码率步一个数量级的事实，为'为什么反向调度更好'提供了非平凡的可量化解释——这是单纯实证研究所缺乏的。第四，**工程上的实用性创新**：提出的路由器 retrofit 方案仅需替换路由而不动 expert 权重，使得已有部署的 LLaDA-MoE 等模型可以低成本享受到 EC 的训练与推理双重收益。
实验结果
本文核心发现分五组实验。**实验一（§3 EC vs TC）**：EC 10.6 小时达 loss 3.75，比 TC dropless (20.7h) 提速约 2.0 倍；EC 达 52.1 TFLOPs/s/GPU，是 TC 的 1.5–2.1 倍；TC per-GPU memory std 3.6 GB，EC 严格相等。**实验二（§4 调度）**：Linear-Reverse 36.5 < Static 37.1 < Cosine-Reverse 37.2 < Gaussian 37.3。**实验三（§4.4 8B-A1B）**：dynamic EC 在 PPL、MMLU、ARC-Challenge 上每 checkpoint 都优于静态 EC。**实验四（§4.5 机理）**：低掩码率 bin $\eta_r$ 高达 $62.2\times10^{-7}$，高掩码率 bin 仅 $9.8\times10^{-7}$，差 7 倍。**实验五（§5 retrofit）**：四任务平均 52.6% → 53.6% / 54.9%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| DLM MoE 训练 (Nemotron-CC, §3) | 训练吞吐 (TFLOPs/s/GPU) | EC 路由: 52.1 | TC dropless: 24.9, TC cap=1.25: 27.0 | 1.5–2.1 倍吞吐提升 |
| DLM MoE 训练 (Nemotron-CC, §3) | 达到 loss 3.75 的 wall-clock 时间 | EC: 10.6 小时 | TC dropless: 20.7h, TC cap=1.25: 19.9h | 约 2.0 倍加速 |
| DLM MoE 推理 (LLaDA-2.0-mini 16B, 8×H100) | per-GPU 显存标准差 (GB) | EC 路由: 0.0 GB (完美均衡) | TC 路由: 3.6 GB (峰值 70.3 GB) | 完全消除负载不均，无 straggler |
| OpenWebText 30B tokens 调度对比 (§4) | 验证困惑度 (PPL) | Linear-Reverse: 36.5 | Static k=20: 37.1, Linear: 37.5, Gaussian: 37.3 | 相对 Static 降低 0.6 PPL（≈1.6% 相对改进） |
| 8B-A1B DLM 预训练 (Nemotron-CC, §4.4) | 验证 PPL @ checkpoint | Dynamic EC (Lin-Rev, k=2–14) 全程优于静态 | Static EC (k=8) | 每个训练 checkpoint 都更优（Figure 5） |
| HumanEval (代码生成, §5 retrofit) | Pass@1 (%) | Dynamic EC: 58.6 | 原 TC: 53.9 | +4.7 个绝对百分点，1.3× 解码加速 |
| HumanEval-Plus (代码生成, §5 retrofit) | Pass@1 (%) | Dynamic EC: 51.6 | 原 TC: 46.1 | +5.5 个绝对百分点 |
| GSM8K (数学推理, §5 retrofit) | 准确率 (%) | Dynamic EC: 73.8 | 原 TC: 74.8 | 略低 -1.0，但在 decode 时间 461s vs 672s 显著更快 |
| MedQA (医学知识, §5 retrofit) | 准确率 (%) | Static EC: 36.7 | 原 TC: 35.6 | +1.1 个百分点 |
| 下游任务平均 (HumanEval/HumanEval+/GSM8K/MedQA, §5) | 平均 Pass@1 / 准确率 (%) | Dynamic EC: 54.9 | TC: 52.6, Static EC: 53.6 | +2.3 个百分点（vs TC），1.3–1.5× 更快解码 |
| 8B-A1B 学习效率机理分析 (§4.5) | 对数损失收敛率 $\eta_r$ (×10⁻⁷) | 低掩码率 bin [0,0.25): 62.2 | 高掩码率 bin [0.75,1.0): 9.8 | 学习效率高出 7 倍以上 |
局限与改进
本文局限性分作者承认与外部观察。作者在 Appendix A 指出，所探索的容量调度是 hand-designed 函数（linear/cosine/Gaussian 及其反向），最优调度可能依赖模型规模、数据集和任务，理论上 learned scheduler（RL policy 或 DiffMoE 风格的 capacity predictor）可能更优。作者还承认 linear-reverse 与 cosine-reverse 在 per-bin $\eta_r$ 上几乎无差别（图 11 ratio 0.97–1.04），但 PPL 仍有 0.7 差距，可能源于'优化轨迹或梯度噪声结构'等高阶效应。外部观察：（1）token drop 风险——附录 C 显示 dynamic EC 在 layer 0 高达 32% drop ratio；（2）retrofit 中 GSM8K 上 dynamic EC (73.8%) 略低于原 TC (74.8%)；（3）评测任务单一；（4）主实验 8B-A1B 是当前规模上限，未在 70B+ 验证。
独立分析的弱点
独立分析本文弱点有四点。**调度函数仍是 hand-designed 且种类有限**：只探索 6 种解析形式，都满足 $\mathbb{E}[s(r)]=\frac{1}{2}$ 对称结构；改进方向是引入 piecewise linear、adaptive threshold 等参数化更强调度族，或训练轻量 capacity predictor。**机理解释停留在相关性层面**：$\eta_r$ 高 7 倍是观察结果，作者仅给出定性解释；改进方向是补充 loss Hessian eigenspectrum 分析。**token drop 隐性风险被低估**：dynamic EC 在 layer 0 的 20–32% drop rate 意味着首层路由尚未成熟时大量 token 只能依赖 shared expert。**retrofit 实验隔离性不足**：Table 6 中 4 种 TC 变体平均准确率 52.2%–53.8%，与 EC (53.6%)、Dynamic EC (54.9%) 差距较小；GSM8K 上 TC (74.8) 反而最高。
未来方向
未来研究方向可分作者提出与可延伸两条线。**作者提出的未来工作**：（1）将 hand-designed 调度替换为 learned 调度，如 end-to-end 训练轻量 capacity predictor（借鉴 DiffMoE）或用 RL policy 自适应容量；（2）更细粒度的 per-token 或 per-block 容量分配；（3）把时间步自适应与深度自适应（early-exit / MoD）正交组合，实现二维自适应计算；（4）探索 EC 在 DLM 推理阶段的极致并行解码。**可延伸方向**：（1）**多模态 DLM 的 EC 路由**：当前 EC-DIT 在视觉扩散已验证，但视觉-语言多模态 DLM 中尚未研究；（2）**KV cache 兼容**：DLM 推理加速（Fast-dLLM 系列）依赖 KV cache，EC 路由下被 expert 选中的 token 集合每步都变，KV cache 复用需重新设计；（3）**更大规模迁移验证**：8B-A1B → 70B+ 量级；（4）**理论分析**：建立 DLM 时间步学习效率与样本复杂度的形式化下界。
复现评估
本文复现性整体良好。**代码与权重**：作者在 GitHub 公开完整代码仓库 https://github.com/zhangshuibai/EC-DLM。**数据**：所有数据均为公开——Nemotron-CC、OpenWebText、GSM8K-AUG-NL、OpenCodeInstruct、MedQA、HumanEval(+)、MMLU、ARC-Challenge。**框架**：基于 Megatron-LM，使用 SwiGLU 激活与 AdamW + WSD 学习率调度。**算力门槛**：附录 D 列出全部超参数——16 层 Transformer、hidden 2048（主实验）或 512（OpenWebText）、64/16/512 fine-grained expert、序列长度 513/2049、global batch 256/288。复现主实验需 H100 GPU 集群（8 卡用于推理显存分析），数天到一周训练时间；retrofit 实验单节点 8 卡即可。**潜在障碍**：附录 E.3 指出吞吐量测量基于前向 FLOPs 公式。
论文图表
