E-MoE: Enhanced Mixture-of-Experts for Non-Factorized Diffusion Language Models
Arseny Ivanov, Alexander Kolesov, Alexander Korotin, Ivan Oseledets, Mikhail Goncharov
cs.CL
2026-09-29
把 MoE 路由决策复用成离散隐变量,扩散语言模型两步内生成困惑度降 2 到 2.6 倍,活跃参数持平
掩码扩散模型(MDM)对 AR 的核心卖点是并行:反向去噪每步一次前向就能揭开多个 token,T 步出全长序列。但它的反向过程按位置因式分解,每个被 mask 的位置各自预测一个边缘分布,同一步揭开的 token 被当作条件独立采样。结果是每个 token 单看合理,拼在一起可能根本不是一句存在过的话。步数越少、每步揭开的 token 越多,这个因式分解误差越重,而少步恰恰是扩散拼速度的地方。
最近的一条线(VADD,Xie et al. 2026)引入连续高斯隐变量、按 VAE 方式联合训练一个识别网络来建模跨位置相关。代价是要预先固定一个高斯先验、训练时多养一个网络,而且容易 posterior collapse:隐变量被解码器静默忽略,信息量趋近于零。
E-MoE 的出发点是:MoE 骨干本来就在每层为每个 token 做 expert 路由决策,这个决策本身就是一个现成的离散隐变量,不必另外发明。
先验与后验共享参数是设计里最关键的一处省略:训练的全部难度收敛到一件事,噪声状态下选的路由要能对齐干净状态下选的路由。LM1B 上活跃参数 139.4M,对 MDLM 的 139.3M,几乎持平。
LM1B(DiT-small,L=128,1M 步,65B token):
| 方法 | Gen-PPL @NFE=1 | @NFE=2 | @NFE=128 |
| MDLM | 1433.8 | 997.5 | 126.3 |
| VADD | 1270.8 | 763.2 | 118.4 |
| E-MoE | 643.8 | 383.7 | 126.3 |
样本熵全部对齐在 4.34–4.35(真实数据 4.32),困惑度增益没有用多样性换。MAUVE 拉得更开:NFE=2 时 E-MoE 24.6%,MDLM 2.95%,VADD 3.76%。步数上去之后优势收窄,NFE≥32 三家 MAUVE 都在 86–92%,E-MoE 反而落后最多 3 个点。AR 参考线(128 步)是 67.97 的困惑度、96.5% 的 MAUVE,少步扩散离它还有明显距离。
其他设置:二值化 MNIST 上 BPD 0.062,好于 VADD 的 0.064 和 MDLM 的 0.077;8-modes 玩具分布在 NFE=1 的 validity 94.4%,MDLM 只有 37.7%。
两个消融把机制坐实了。同一个 checkpoint 把采样路由从随机改成贪心(温度趋零,混合塌缩成单 expert),NFE=1 困惑度回升 1.72 倍、NFE=2 回升 2.14 倍,朝 MDLM 方向退回去;权重一个没动,增益确实来自随机混合本身。训练动态上,VADD 的隐变量 KL 掉到 0.14 nats/token,接近塌缩;E-MoE 的路由 KL 稳在 0.65,全程没有 warmup 和辅助损失。
扩散语言模型要跟 AR 拼解码延迟,少步质量是主战场。E-MoE 在架构内部解决:不加识别网络、不设先验、活跃参数持平,训练管线和稠密基线一样简单。对 LLaDA-MoE 这类骨干本来就是 MoE 的扩散模型,路由决策是现成的,改的只是训练目标,迁移成本很低。这是机制层面的改动,不是蒸馏那种需要先训好 teacher 的后处理。
作者自己列了两条:高 NFE 下每步揭开的 token 变少,收益自然衰减;整个机制押在路由对齐上,噪声 router 恢复不了干净路由、或路由分布塌缩,混合就退化回单个因式分解分量。另外几点论文没有覆盖:验证规模停在 LM1B、无条件生成、139M 活跃参数,离 LLaDA 级别还差近两个数量级;总参数 536M 对 139M,显存与部署并不免费;MNIST 上需要 load-balancing 损失、LM1B 上又不需要,KL 权重和温度调度在不同任务间不统一,超参敏感性没讨论;MAUVE 在 NFE=64 之后输给对手一到三个点,「少步之外不占优」正文提得不多。