[ICLR 2025] MDM-Prime-v2：扩散模型首次在缩放效率上击败自回归架构

TL;DR

长期以来，自回归模型（ARM）凭借其极高的计算效率统治着 LLM 领域，而扩散语言模型（MDM）虽然具备并行生成的潜力，却一直受困于高达 16 倍的效率赤字。本文提出的 MDM-Prime-v2 通过二进制编码和索引打乱两大核心手段，改写了这一局面。实验证明，它不仅在 Perplexity 上全面碾压同规模 ARM，更在计算效率上实现了 21.8 倍的阶跃。

背景定位：打破“自回归霸权”

在学术界，离散扩散模型被认为是通往非线性、非顺序文本生成的关键方案。然而，要在扩展性（Scaling）上与 GPT 系列竞争，MDM 必须解决其变分界（Variational Bound）不够紧致的问题。MDM-Prime-v2 的出现，标志着扩散模型正式进入了“计算优化（Compute-optimal）”的竞争序列。

痛点深挖：为什么以前的扩散模型“不给力”？

作者通过数学推导发现，MDM 的性能高度依赖于**子标记化（Subtokenization）**的设计。

1. 粒度玄学：以前选取标记粒度（Granularity）多靠经验性尝试，缺乏数学上的单调性证明。
2. 熵值塌陷：常用的 BPE 分词器生成的索引具有极强的结构性（频率越高索引越小）。直接进行进制转换会导致子标记的物理含义极度贫乏（熵极低），模型在去噪时无法获得足够明确的语义指导。

核心算法：二进制编码与索引重排

为了收紧变分下界，MDM-Prime-v2 实施了手术般的改进：

1. 二进制编码（Binary Encoding）

理论推导证明（Proposition 3.1 & 3.2），变分界随粒度 $\ell$ 的增加而收紧。作者索性将 $\ell$ 推至极限——二进制编码。这意味着模型是在比特位级进行采样和去噪，提供了最精细的预测粒度。

2. 索引打乱（Index Shuffling）

这是本文最具直觉的贡献。 BPE 索引分布极不均匀，通过在进入嵌入层前进行一次随机索引打乱（Lookup Table 形式，零额外计算开销），可以人工“制造”出高熵的子标记分布。如下表所示，打乱后的熵接近理论最大值：

实验与结果：算力效率的降维打击

在 Scaling Law 的视角下，MDM-Prime-v2 展示了惊人的潜力：

• Perplexity 突破：在 OpenWebText 上，MDM-Prime-v2 达到了 7.77 PPL，而同预算的自回归模型仅为 12.99。
• 计算效率：实现相同的 Loss 级别，MDM-Prime-v2 所需的 FLOPs 远少于 ARM。

此外，在 1.1B 参数规模的 Zero-shot 任务中，该模型在 SciQ 和 McTaco 等理性和时间推理任务上表现尤为卓越，证明了其表征能力的深度。

深度洞察：为什么这种做法有效？

通过对注意力矩阵的分析（Figure A10），作者发现 MDM-Prime-v2 完全规避了自回归模型中常见的 Attention Sink（注意力沉没） 现象。由于引入了更复杂的子标记交互，其注意力模式呈现出丰富的对角线结构和多样性，这暗示模型学会了更高级的路由机制。

总结与启示

MDM-Prime-v2 告诉我们：

1. 底层编码决定上限：标记化的数学特性直接限制了扩散模型的似然估计潜力。
2. 非自回归大有可为：只要解决了编码效率问题，扩散模型在计算效率上具备超越 Transformer-ARM 的可能性。

局限性分析：尽管在预训练指标上表现优异，但扩散模型在极高掩码率（$t\to 1$）时的联合分布建模仍有提升空间，这将是未来从 Pretraining 跨向更强 Alignment 的关键。