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

📄 arXiv: 2603.16077v1 📥 PDF

作者: Chen-Hao Chao, Wei-Fang Sun, Junwei Qua, Chun-Yi Lee, Rahul G. Krishnan

分类: cs.LG

发布日期: 2026-03-17

备注: Project Page: https://chen-hao-chao.github.io/mdm-prime-v2

💡 一句话要点

MDM-Prime-v2：通过二元编码和索引混洗实现扩散语言模型计算最优扩展

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 扩散语言模型 掩码模型 二元编码 索引混洗 计算效率 自然语言生成 零样本学习

📋 核心要点

  1. 现有MDM-Prime框架在token粒度超参数选择上缺乏有效指导，且与BPE tokenizer结合时似然估计显著降低。
  2. 论文提出MDM-Prime-v2，通过引入二元编码和索引混洗，优化变分下界，提升模型性能和计算效率。
  3. 实验表明，MDM-Prime-v2在计算效率上优于自回归模型，并在OpenWebText和常识推理任务上取得显著提升。

📝 摘要（中文）

掩码扩散模型(MDM)在使用Partial masking scheme (Prime)学习时表现出卓越的泛化能力。这种方法将token转换为sub-token，并在sub-token级别对扩散过程进行建模。我们发现了MDM-Prime框架的两个局限性。首先，我们缺乏指导sub-tokenizer中token粒度超参数选择的工具。其次，我们发现当sub-tokenizer与常用的Byte-Pair-Encoding (BPE) tokenizer配对使用时，sub-tokenizer的函数形式会显著降低似然估计。为了解决这些局限性，我们研究了MDM-Prime中变分下界的紧密性，并开发了MDM-Prime-v2，这是一种结合了二元编码和索引混洗的掩码扩散语言模型。我们的扩展分析表明，MDM-Prime-v2的计算效率是自回归模型(ARM)的21.8倍。在计算最优的比较中，MDM-Prime-v2在OpenWebText上实现了7.77的困惑度，优于ARM (12.99)、MDM (18.94)和MDM-Prime (13.41)。当将模型大小扩展到1.1B参数时，我们的模型在各种常识推理任务上进一步展示了卓越的零样本准确率。

🔬 方法详解

问题定义：现有的Masked Diffusion Model (MDM) Prime框架存在两个主要问题。一是缺乏有效的方法来指导sub-tokenizer中token粒度超参数的选择，导致难以优化模型性能。二是当MDM-Prime与常用的Byte-Pair-Encoding (BPE) tokenizer结合使用时，sub-tokenizer的函数形式会显著降低似然估计，影响模型的生成质量。

核心思路：论文的核心思路是通过引入二元编码和索引混洗来改进MDM-Prime框架。二元编码旨在更有效地表示sub-token，减少信息冗余，从而提高模型的学习效率。索引混洗则通过改变sub-token的排列顺序，增加模型学习的难度，从而提高模型的泛化能力。

技术框架：MDM-Prime-v2的整体框架仍然基于Masked Diffusion Model，但引入了两个关键模块。首先，在token化阶段，使用二元编码器将token转换为二元sub-token序列。其次，在扩散过程之前，对sub-token序列进行索引混洗，改变其排列顺序。然后，模型在混洗后的sub-token序列上进行扩散和逆扩散过程，最终生成文本。

关键创新：MDM-Prime-v2最重要的技术创新点在于二元编码和索引混洗的结合使用。二元编码能够更有效地表示信息，减少冗余，提高学习效率。索引混洗则能够增加模型学习的难度，提高模型的泛化能力。这种结合使得MDM-Prime-v2在计算效率和模型性能上都优于现有的MDM和MDM-Prime模型。

关键设计：二元编码器的具体实现方式未知，但其目标是将token转换为二元序列。索引混洗的具体实现方式也未知，但其目标是随机改变sub-token的排列顺序。损失函数仍然基于变分下界，但由于二元编码和索引混洗的引入，其形式可能有所改变。具体的网络结构可能与MDM-Prime类似，但需要进行相应的调整以适应二元sub-token序列的输入。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MDM-Prime-v2在OpenWebText数据集上取得了显著的性能提升，困惑度达到7.77，远优于ARM (12.99)、MDM (18.94)和MDM-Prime (13.41)。此外，在扩展到1.1B参数的模型规模后，MDM-Prime-v2在各种常识推理任务上表现出卓越的零样本准确率，证明了其强大的泛化能力。

🎯 应用场景

MDM-Prime-v2具有广泛的应用前景，可应用于自然语言生成、文本摘要、机器翻译等领域。其高效的计算性能使其更适用于资源受限的场景。通过进一步优化，该模型有望在各种常识推理和知识密集型任务中发挥重要作用，推动人工智能技术的发展。

📄 摘要（原文）

Masked diffusion models (MDM) exhibit superior generalization when learned using a Partial masking scheme (Prime). This approach converts tokens into sub-tokens and models the diffusion process at the sub-token level. We identify two limitations of the MDM-Prime framework. First, we lack tools to guide the hyperparameter choice of the token granularity in the subtokenizer. Second, we find that the function form of the subtokenizer significantly degrades likelihood estimation when paired with commonly used Byte-Pair-Encoding (BPE) tokenizers. To address these limitations, we study the tightness of the variational bound in MDM-Prime and develop MDM-Prime-v2, a masked diffusion language model which incorporates Binary Encoding and Index Shuffling. Our scaling analysis reveals that MDM-Prime-v2 is 21.8$\times$ more compute-efficient than autoregressive models (ARM). In compute-optimal comparisons, MDM-Prime-v2 achieves 7.77 perplexity on OpenWebText, outperforming ARM (12.99), MDM (18.94), and MDM-Prime (13.41). When extending the model size to 1.1B parameters, our model further demonstrates superior zero-shot accuracy on various commonsense reasoning tasks.