并非所有去噪步骤都等价：用模型调度加速掩码扩散语言模型 Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models
让小模型接管扩散轨迹首尾步骤，中间保留大模型，省 17% 算力几乎不掉质量
前置知识
掩码扩散语言模型 (MDLM)
MDLM 是一类离散扩散语言模型。前向过程按时间 $t\in[0,1]$ 以概率 $(1-\alpha_t)$ 把 token 替换成特殊掩码符号 $m$（本文用线性 $\alpha_t=1-t$，期望掩码比例正好等于 $t$）。反向过程从全掩码序列出发，用一个双向 Transformer 反复预测被掩码位置的原 token，训练目标是只在掩码位置上的加权交叉熵 $\mathcal{L}_{\mathrm{MDLM}}$。与自回归逐 token 生成不同，它每步可并行更新多个 token，且一旦某 token 被解掩码就固定不变。
本文全部方法都建立在 MDLM 的『多步全序列去噪』结构上——正是这个结构带来了昂贵的采样成本，也提供了『在不同步骤换用不同容量模型』的操作空间。
KV 缓存 (KV Cache) 的缺失
自回归模型解码时可以缓存已生成 token 的 Key/Value，使每步只需计算新 token，成本很低。但 MDLM 是双向注意力、每步都要对整条序列做一次完整前向传播，token 之间的可见关系随去噪不断变化，因此无法直接复用 KV 缓存。
这是 MDLM 采样昂贵、成为实际瓶颈的根本原因，也是本文寻求『用小模型省算力』这一正交加速思路的动机所在。
生成困惑度 (Generative Perplexity)
衡量无条件生成质量的指标：先用被评估模型采样出文本，再用一个独立的预训练模型（本文用 GPT-2）计算这些样本的困惑度，越低表示生成文本越『通顺/像自然语言』。本文默认生成 1600 条长度 1024 的样本、$T=1000$ 步来求均值。作者也承认该指标在某些情况下不可靠，但在统一训练与采样协议下用作相对比较是可靠的。
它是本文比较不同调度方案质量优劣的核心度量，理解其含义和局限才能正确解读所有实验结论。
步骤重要性 / 时间步依赖的算力分配
扩散生成经过一串时间步，从高噪声逐渐变清晰。图像扩散领域已发现不同时间步对模型容量的『需求』不同，常呈平滑或单调趋势，于是有 T-Stitch、DyDiT、OMS-DPM 等方法在部分步骤跳过/缓存/换小模型。本文要研究的正是：文本掩码扩散是否也有这种非均匀性、其模式是否和图像一致。
本文的核心发现（中间步骤最敏感、首尾步骤最鲁棒，呈『峰型』而非图像扩散的单调型）正是对这一概念在离散文本扩散上的首次系统刻画。
研究动机
MDLM 虽然在质量上已逼近自回归语言模型，但采样代价高昂：生成一条序列需要多达 $T=1000$ 次对整条序列的全量 Transformer 前向传播，而且由于是双向注意力、每步可见关系不断变化，无法像自回归解码那样借助 KV 缓存降低成本。结果是即便 MDLM 的质量很强，推理成本仍是实际部署的瓶颈。作者在 profiling 中还指出，对 4-block 小模型而言输出词表投影层竟占据约 81.6% 的运行时，Transformer 块只占约 18.2%，说明成本结构复杂，单纯堆步数或大模型都很浪费。这构成了『如何在几乎不损失质量的前提下降低单条样本采样算力』这一核心问题。
本文的目标是本文目标是提出一种简单、与架构无关、且无需重训重模型/蒸馏/改采样算法的推理期加速策略——模型调度 (model scheduling)：在去噪轨迹的部分时间步上，用一个单独训练好的更小 (light) MDLM 替换大 (heavy) MDLM。核心要回答的问题是：哪些时间步对模型替换最鲁棒？light 与 heavy 步骤应如何排布，才能在速度与质量之间取得最优折衷？作者希望给出一条『简单、可解释、跨数据集通用』的调度经验法则，并从机理上解释为什么某些步骤更敏感。
与已有工作不同的是，已有的时间步依赖算力分配方法（T-Stitch、OMS-DPM、DyDiT、MD-DiT、Flexiffusion 等）几乎都在连续图像扩散上开发和验证，其结论往往是『步骤难度随轨迹平滑/单调变化』。但文本掩码扩散在状态空间（离散 token + 掩码）和预测结构（词表上的类别分布、不确定性集中在掩码位置）上都与图像扩散本质不同，因此这些结论未必能迁移。本文的独特切入点，就是首次在离散掩码文本扩散上系统研究步骤重要性，并发现了与图像扩散单调趋势截然不同的『中间峰型』结构。此外它与减少去噪步数、恢复 KV 缓存类的加速方法正交，理论上可叠加得到乘性加速。
核心方法
直觉上，扩散生成的每一步『难度』不同：轨迹两端（接近全掩码或接近全清晰）也许简单到用小模型就能胜任，中间某些步则可能是决定生成质量的关键。本文把这一直觉落地为『模型调度』：训练一族只在深度上不同（4/6/8/10/12 个 Transformer 块，宽度固定 hidden 768、MLP ratio 4、同词表同 tokenizer）的 MDLM，12-block 作 heavy 基线，更小的作为候选 light 模型。定义调度函数 $s:\{1,\dots,T\}\to K$，在反向第 $i$ 步（$t_i=i/T$）选用 $p_{\theta_{s(i)}}(\cdot\mid z_{t_i},t_i)$ 采样。若 heavy 有 $B_H=12$ 块、light 有 $B_L$ 块，替换比例 $p$ 的步骤，可省算力约 $\text{saved FLOPs}\approx p\cdot\frac{B_H-B_L}{B_H}$。作者通过手工设计调度、穷举搜索、以及跨时间步的相似度分析三条线索，共同确认最优排布规律。
核心创新是把『在哪些步骤上花算力』当作一个可自由设计的调度问题，并用实验揭示离散文本扩散的步骤重要性呈『中间峰型』：固定算力预算时，light 步骤放在轨迹中段（2/3 象限）质量最差，放在两端最好。由此提炼出 sandwich（三明治）调度——把 light 步骤对半劈开放到轨迹首尾，如 (L125, H750, L125)。这与自回归 T-Stitch 那种『只把小模型塞到轨迹前段』不同，也与图像扩散报告的单调趋势本质不同。更关键的是，该方法不改采样器、不蒸馏、不改架构，仅仅换『每步跑哪个模型』，因而与减步数、KV 缓存类方法正交，可组合叠加。
方法步骤详情
第一步：训练一族仅深度不同的 encoder-denoiser（4/6/8/10/12 块），沿用 MDLM (Sahoo et al., 2024) 的代码与默认设置，在 OpenWebText 上用 GPT-2 tokenizer 训 1M 步、有效 batch 512、序列长 1024（约 262B 掩码 token），AdamW、lr $3\times10^{-4}$、2500 步 warmup、$\beta=(0.9,0.999)$；并在 LM1B（128 token）上重复训练同族模型验证通用性。第二步：固定 light 步比例（先 25%，即 250/1000 步，对 4-block 省 16.7% FLOPs），手工构造把 250 个 light 步放在不同象限的调度，以及 sandwich 调度，比较生成困惑度。第三步：为排除手工调度偏差，做穷举搜索——把 1000 步切成 10 个连续段（各 100 步），选 4 段跑 light（40% 步、省 26.7% FLOPs），共 $\binom{10}{4}=210$ 种，每种用 160 条固定种子样本评估。第四步：扫描两个尺度维度——固定 sandwich 排布变 light 深度（4→10 块），以及固定 (12-block, 4-block) 变 light 步比例（0%→100%）并测实际墙钟时间。第五步：机理分析——在每个固定时间步、对相同被污染输入 $z_t$ 比较 light 与 heavy 的掩码交叉熵绝对差 $\Delta_{\mathrm{loss}}(t)$ 与 token 级 KL 散度 $\Delta_{\mathrm{KL}}(t)$（并减去两个不同种子 heavy 之间的基线 KL）。第六步：把穷举搜索结果转成段影响力分数 $I(j)$。第七步：在 prefix-conditional（256/128 token 前缀）设定下重复验证，并用 token 级熵检查多样性。
技术新颖性
技术新颖性有三点。其一，首次把『模型调度』这一图像扩散里的思路系统地搬到离散掩码文本扩散，并证明它可行、可解释。其二，发现并用三种独立证据（手工调度、$\binom{10}{4}=210$ 穷举搜索、loss/KL 相似度曲线）交叉验证了『中间步骤最敏感』的峰型结构——$\Delta_{\mathrm{loss}}$ 和相对 KL 都在 $t\approx0.4$–$0.6$ 处出现明显峰值，段影响力分数 $I(j)$ 也显示中间段为正（替换有害）、首尾段为负（替换安全，最极端两端分别达 $-0.80$ 与 $-0.44$）。这与图像扩散的单调趋势形成鲜明对照，是对领域认知的实质性修正。其三，方法完全免训练、架构无关，且明确论证了它与减步数、KV 缓存类加速的正交性与可组合性，具有很好的工程实用价值。
实验结果
核心发现是掩码文本扩散的步骤重要性强烈非均匀，且呈中间峰型。(1) 固定 25% light 步（4-block，省 16.7% FLOPs）时，把 light 步放中段困惑度最差，sandwich (L125,H750,L125) 最好，达到 44.3，紧随其后是全放第一象限，而全 heavy 基线为 42.9，即仅 3.4% 相对退化换来 16.7% 算力节省。(2) Table 1 显示固定 sandwich 排布、light 深度从 4→10 块，生成困惑度分别为 44.31/43.67/43.45/42.90，相对 12-block 基线 (42.85) 的退化从 3.41% 降到 0.12%，对应省 FLOPs 从 16.67% 降到 4.17%——深度越大质量越好但省得越少，构成连续折衷。(3) Table 2 固定 (12b,4b) 把 light 步比例从 0%→100%，困惑度从 42.9 平滑升到 53.4，省 FLOPs 从 0%→66.7%，实测墙钟加速从 0%→28.3%，说明这是连续的速度–质量旋钮。(4) 穷举搜索（Figure 2/3/6）中最优的 top-5 都把 light 段放两端、最差 bottom-5 都放中间，段频率与段影响力分数一致印证峰型。(5) 机理分析 Figure 4/5：light 与 heavy 的 loss 绝对差与相对 KL 都在轨迹中部 $t\approx0.4$–$0.6$ 出现峰值，为『中间替换最有害』给出机制解释。(6) Table 3 前缀条件生成（256 token 前缀）排名不变，sandwich 40.9 对全 heavy 39.3，且所有调度 token 级熵稳定在 5.40–5.42，多样性无损。(7) 上述模式在 LM1B（128 token）上完整复现（Figure 11/12、Table 4），跨数据集通用。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| OpenWebText 无条件生成 (25% light, 4-block) | 生成困惑度 (GPT-2) | sandwich 44.31 ± 0.76 | 全 12-block heavy 42.85 ± 0.71 | 仅 +3.41% 困惑度退化，省 16.67% FLOPs |
| OpenWebText 无条件生成 (25% light, 10-block) | 生成困惑度 (GPT-2) | 42.90 ± 0.70 | 42.85 ± 0.71 | +0.12% 退化，省 4.17% FLOPs |
| OpenWebText 无条件 (100% light 4-block) | 生成困惑度 / 墙钟加速 | 53.4 / 28.3% 加速 | 42.9 / 0% (全 heavy) | 省 66.7% FLOPs，速度–质量连续折衷极端点 |
| OpenWebText 前缀条件生成 (256 token 前缀) | 条件生成困惑度 / 熵 | sandwich 40.9 ± 0.36 / 熵 5.42 | 全 heavy 39.3 ± 0.37 / 熵 5.40 | 排名不变，多样性无损 (熵 5.40–5.42) |
| LM1B 无条件生成 (32/128 light, 4-block) | 生成困惑度 / 熵 | sandwich 44.3 ± 0.50 / 熵 5.27 | 全 heavy 42.9 ± 0.47 / 熵 5.29 | 跨数据集复现峰型，熵变化 <0.03 nats |
局限与改进
作者坦承的局限：其一，评估主要依赖生成困惑度，而该指标本身在某些情形下可能不可靠（ReMDM 讨论过失效模式），作者仅将其作为统一协议下的相对度量。其二，最关键的是 FLOPs 节省与实际墙钟加速存在明显落差：由于输出词表投影层在 4-block 模型上占约 81.6% 运行时（12-block 上仍占约 59.9%），减深度只影响约 18.2%–40% 的块算力，因此 Table 2 中 66.7% 的 FLOPs 节省只换来 28.3% 的实测加速。作者建立了简单模型 $\text{speedup}\le \frac{1}{1-\alpha}$（$\alpha=T_{\mathrm{out}}/T$）解释上限，并指出这非本质瓶颈，用 Liger-Kernel、TensorRT-LLM、FlashInfer 等融合/优化核可缩小差距。其三，实验规模有限：heavy 仅 12 块、hidden 768，缺少大规模预训练 MDLM 族的验证。我的补充观察：穷举搜索限定在 10 个连续段的粗粒度离散空间，可能错过更细粒度的最优调度；且只对比了两模型两两组合，多模型/动态路由尚未探索。
独立分析的弱点
弱点一：墙钟收益受输出层主导严重稀释——在 hidden 768 的小模型上词表投影占八成运行时，导致模型调度的实际加速远低于 FLOPs 预测。改进方向是先用融合投影–softmax 核（Liger-Kernel）或高效推理栈（TensorRT-LLM/FlashInfer）优化非深度相关部分，再叠加调度，或在更大 hidden/更长序列（块算力占主导）的场景下应用。弱点二：调度是静态的、按固定象限硬编码，无法根据当前去噪状态自适应。改进方向是引入 early-exit 或以去噪状态为条件的路由策略，让每步动态决定容量。弱点三：只用两个模型（12b + 4b），中间敏感段被迫用最大模型、两端被迫用最小模型，粒度粗。改进方向是扩展到多容量层级（4/6/8/10/12 混用），在敏感度曲线上做更平滑的容量匹配。弱点四：穷举搜索受限于 10 段离散化，且靠生成困惑度这一可能不可靠的指标做排序，结论稳健性依赖该指标。改进方向是用多种质量指标（MAUVE、下游任务表现）交叉验证，并在更细段或连续参数化上做可微搜索。
未来方向
作者明确提出三条：(1) 当出现像 Qwen、LLaMA 那样成体系、多尺度的预训练 MDLM 族时，应在更大规模、标准 benchmark 上验证本文发现；(2) 把调度从两模型推广到多容量层级；(3) 引入 early-exit 或以去噪状态为条件的动态路由，进一步改善速度–质量折衷。基于成果可延伸的方向还有：把模型调度与减步数解码器（FlashDLM、LocalLeap、CD4LM）及 KV 缓存类方法（Fast-dLLM、dInfer）组合，验证作者所说的乘性加速；把中间峰型的机理分析进一步与『掩码位置不确定性集中』联系起来，找出为什么中间步骤 $t\approx0.5$ 信息量最大；以及探索让 light/heavy 的边界随 prompt 难度自适应移动。
复现评估
复现友好度总体较高。方法本身极简，仅需训练一族只在深度不同、其余完全相同的 MDLM，并严格沿用了 MDLM (Sahoo et al., 2024) 的公开代码库与默认设计，训练超参数完整给出（OpenWebText、GPT-2 tokenizer、1M 步、有效 batch 512、序列 1024、约 262B 掩码 token、AdamW、lr $3\times10^{-4}$、2500 warmup、$\beta=(0.9,0.999)$）。评测协议明确（GPT-2 生成困惑度、1600 条长 1024 样本、$T=1000$；LM1B 用 128 token）。算力门槛：需从零训练 5 个规模的模型各 1M 步，属中等偏上成本，但单模型只有 12 块、hidden 768，学术实验室可承受。穷举搜索 210 种调度每种 160 样本也可控。主要不确定性是论文（预印本）未明确说明是否开源自己的训练脚本与 checkpoint——若仅依赖上游 MDLM 代码库需自行补齐模型族训练与调度评测逻辑，但由于设计透明、超参完整，独立复现难度中等。
论文图表
