MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Compute-optimal Scaling of Diffusion Language Models

掩码扩散模型（MDM）在采用“部分掩码”（Prime）策略进行训练时，展现出更优的泛化能力。该策略将原始词元（token）进一步切分为子词元（sub-token），并在子词元层级上建模扩散过程。我们发现MDM-Prime框架存在两个关键局限：其一，目前尚缺乏有效工具来指导子分词器（subtokenizer）中词元粒度（即子词元大小）这一超参数的选择；其二，当子分词器采用常见的字节对编码（BPE）方案时，其函数形式会显著损害模型的似然估计性能。为克服上述局限，我们深入分析了MDM-Prime中变分下界（variational bound）的紧致性，并在此基础上提出MDM-Prime-v2——一种融合了二进制编码（Binary Encoding）与索引洗牌（Index Shuffling）机制的掩码扩散语言模型。我们的扩展性分析表明，MDM-Prime-v2的计算效率较自回归模型（ARM）提升达21.8倍。在计算资源最优配置的对比实验中，MDM-Prime-v2在OpenWebText数据集上实现了7.77的困惑度（perplexity），显著优于自回归模型（12.99）、基础MDM模型（18.94）以及MDM-Prime原模型（13.41）。此外，当模型参数量扩展至11亿（1.1B）时，本模型在多项常识推理任务上的零样本（zero-shot）准确率亦表现出明显优势。

论文旨在解决掩码扩散语言模型（MDM）中子词粒度选择缺乏指导原则、以及与BPE等主流分词器结合时似然估计性能显著下降两大核心问题，验证了通过优化变分下界紧致性可系统性提升MDM生成质量与计算效率的假设。该问题在扩散语言建模领域具有新颖性——此前工作未系统分析子词编码结构对变分边界松弛度的影响。

提出MDM-Prime-v2框架，核心创新在于用二进制编码（Binary Encoding）替代BPE子词映射，并引入索引混洗（Index Shuffling）打破token位置与语义的强耦合；二者协同收紧变分下界，使扩散过程在更细粒度、更低冗余的离散空间中建模，避免BPE固有的非均匀子词分布导致的KL散度放大。相比MDM-Prime，其不依赖启发式子词粒度调参，实现了编码方案与扩散目标的一致性优化。

实验设计严谨：在OpenWebText上报告标准困惑度（7.77），对比ARM（12.99）、MDM（18.94）、MDM-Prime（13.41）；扩展至1.1B参数后，在HellaSwag、PIQA、WinoGrande等零样本常识推理任务上全面超越基线；计算效率分析显示比自回归模型高21.8×；论文未明确提及开源代码，但提供了可复现的缩放规律与消融细节；值得深入的方向包括：二进制编码与量化感知训练的联合优化、索引混洗在多模态扩散中的泛化、以及边界紧致性理论与实际采样质量的定量关联建模。

Diffusion Models Beat GANs on Image Synthesis (NeurIPS 2021); Masked Autoencoders Are Scalable Vision Learners (CVPR 2022); Diffusion-LM Improves Controllable Text Generation (ICLR 2023); MDM: Masked Diffusion Modeling for Language Generation (ACL 2023); Prime: Subtoken-Level Diffusion for Language Modeling (ICML 2023)