Dependency-Guided Parallel Decoding in Discrete Diffusion Language Models¶
会议: NeurIPS 2026
arXiv: 2604.02560
代码: https://github.com/liranringel/DEMASK
领域: LLM 其他
关键词: 离散扩散语言模型, 并行解码, 依赖关系预测, 非自回归生成, 总变差距离界限
一句话总结¶
针对离散扩散模型并行反掩码时因强耦合 token 引发的边际-联合分布失配问题，本文提出轻量级依赖预测器 DEMASK，仅需单次前向传播即可预估掩码位置间的条件依赖矩阵，并通过贪心选择在全变差距离理论保证下实现 1.7–2.2× 加速且维持甚至提升精度。
研究背景与动机¶
离散扩散语言模型（dLLM）通过迭代去噪逐步反掩码生成文本，相较于传统自回归模型严格从左至右的因果依赖，具备双向上下文建模能力，并且天然支持在单个扩散步骤内并行解掩多个 token，从而显著降低采样步数。然而，现有的大规模 dLLM 在并行解码时面临一个根本性的分布失配难题：模型的真实联合条件分布遵循链式法则，而并行解码在选定位置子集上只能独立抽取各位置的边际分布乘积。当被选中的 token 之间存在强依赖时，这种将联合分布近似为完全可分离边际乘积的假设会导致严重的语义不一致，最典型的例子如计算表达式“1 + □ = □”时，两个空位各自的边际可能均偏好数字“2”或“3”，但独立并行采样却极易组合出“1 + 2 = 5”的荒谬输出。
目前现有的主流解码策略主要依赖置信度筛选（如基于熵或 Top-1 概率）、步间分布稳定性（如 KLASS 利用步间预测分布的 KL 散度）或重掩码调度。这些启发式策略将每个 token 视为孤立个体，没有显式度量被选位置之间的条件相互影响；而另一些尝试建模依赖的工作（如 PUNT 或辅助自回归混合模型）要么在每个时间步需要引入 \(\mathcal{O}(\log |M|)\) 次前向计算开销，要么依赖额外的辅助自回归模型。如何在不显著增加计算开销的前提下，精准识别在当前上下文下近似条件独立的掩码位置子集，成为打通扩散语言模型高速与高质量生成的关键瓶颈。
本文的切入视角在于：只要能够在单次前向传播中快速预估任意两个掩码位置之间的成对条件影响，就能通过受控的贪心累加策略找出累积依赖有界的子集并行反掩码。核心 idea：在扩散骨干网络顶层外挂一个轻量级依赖预测器，仅利用单次前向隐状态预估掩码 token 间的成对全变差（TV）依赖矩阵，并结合置信度阈值与贪心算法选出累积依赖低于阈值 \(\tau\) 的近似条件独立集合，在次可加性假设下为并行采样与真实联合分布的总变差距离提供严格理论上界。
方法详解¶
整体框架¶
DEMASK 采用迭代扩散解码循环，每个去噪步骤主要由单次前向特征提取、贪心子集选择以及并行采样更新三个阶段串联构成。输入为 prompt 上下文 \(X\) 以及当前已解掩 token 集合 \(Y_U\)，输出为逐步填充完毕的完整序列。骨干网络在执行常规去噪前向传播的同时，输出顶层隐状态；外挂的依赖预测器将隐状态映射为成对依赖矩阵 \(\hat{\mathbf{D}}\)，随后贪心子集选择算法结合 Top-1 置信度筛选出满足累积依赖约束的最大并行位置集合 \(S\)，最后模型在集合 \(S\) 上并行抽取 token 并合并进上下文，直到所有掩码位置全部消除。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入上下文与掩码序列<br/>X, Y_U, M"] --> B["前向传播与隐状态提取<br/>获得词表分布 P_hat 与隐状态 H"]
B --> C["双线性依赖预测器<br/>计算成对全变差矩阵 D_hat"]
C --> D["贪心位置子集选择<br/>置信度筛选与累积依赖阈值约束"]
D --> E["并行采样与状态更新<br/>在选定位置 S 独立采样并更新 Y_U"]
E -->|仍有掩码位置| B
E -->|M 为空| F["输出完整生成文本"]关键设计¶
1. 双线性依赖预测器：以单次前向预估掩码位置间的全变差依赖矩阵
为避免对 \(|M|\) 个掩码位置分别做前向探查带来的 \(\mathcal{O}(|M|)\) 巨大开销，本文设计了一个类似缩放点积注意力结构的轻量预测头。真实成对依赖度量定义为当位置 \(m_j\) 被揭示时，位置 \(m_i\) 条件分布的期望全变差（Total Variation, TV）距离： $\(D_{i,j} = \mathbb{E}_{Y_{m_j} \sim P_\theta(\cdot \mid X, Y_U)} \left[ \text{TV}\left( P_\theta(Y_{m_i} \mid X, Y_U), P_\theta(Y_{m_i} \mid X, Y_U, Y_{m_j}) \right) \right]\)$ 预测器直接挂载在冻结骨干网络的最后一层隐藏状态 \(\mathbf{H} \in \mathbb{R}^{|M| \times d}\) 上。在推理阶段，可将学习到的投影矩阵 \(\mathbf{W}_Q, \mathbf{W}_K \in \mathbb{R}^{d \times d}\) 离线融合成单一权重 \(\mathbf{W} = \mathbf{W}_Q \mathbf{W}_K^\top\)，通过单次矩阵乘法直接计算预测依赖矩阵元素： $\(\hat{D}_{i,j} = \sigma\left( \frac{\mathbf{h}_i \mathbf{W} \mathbf{h}_j^\top}{\sqrt{d}} \right)\)$ 其中 \(\sigma(\cdot)\) 为 Sigmoid 函数，严格保证预测输出落在 \([0, 1]\) 之间且对角线置 0。该模块参数量仅为 \(2d^2\)（以 7B 模型的 \(d=3584\) 计算约为 26M 参数），相比骨干网络完全可忽略。
2. 贪心位置子集选择：平衡从左至右先验、置信度过滤与累积依赖上界
在获得预测矩阵 \(\hat{\mathbf{D}}\) 与各位置 Top-1 概率 \(p_{\text{top1}}\) 后，贪心算法首先选入最左侧掩码位置 \(s_1 = \min(M)\) 作为初始种子，这一设计顺应了自然语言局部由前向后的强逻辑展开惯性。随后算法在未选掩码集合中维护一个置信度候选池 \(\mathcal{C} = \{c \in M \setminus S \mid p_{\text{top1}}[c] > \gamma\}\)，防止模型过早敲定高度不确定的位置。在候选池非空的前提下，算法每轮贪心选取对当前已选集合 \(S\) 增量依赖最小的候选位置 \(c^*\)： $\(c^* = \arg\min_{c \in \mathcal{C}} \sum_{s \in S} \hat{D}_{\pi(c), \pi(s)}\)$ 当累积增量 \(\mathcal{D}_{\text{accum}} + \sum_{s \in S} \hat{D}_{\pi(c^*), \pi(s)} \le \tau\) 满足预设上界时，将 \(c^*\) 纳入 \(S\) 并累加依赖；一旦超出阈值 \(\tau\) 或候选池耗尽则立即终止当前步的选择，转入并行采样。
3. 次可加性假设下的分布距离理论界限：全变差距离的可控性保证
传统并行采样使用完全分解乘积 \(Q_\theta(Y_S \mid X, Y_U) = \prod_{i \in S} P_\theta(Y_i \mid X, Y_U)\) 近似真实的链式联合分布 \(P_\theta(Y_S \mid X, Y_U)\)。本文基于次可加性假设（即某变量受多历史变量联合影响的 TV 偏差不超过其与各历史变量成对 TV 偏差之和）： $\(\mathbb{E}_{Y_{s_{<t}}}\left[ \text{TV}\left( P_\theta(Y_{s_t} \mid X, Y_U), P_\theta(Y_{s_t} \mid X, Y_U, Y_{s_{<t}}) \right) \right] \le \sum_{j=1}^{t-1} D_{\pi(s_t), \pi(s_j)}\)$ 通过构造插值分布序列并利用全变差距离的伸缩和与三角不等式，证明了只要贪心算法控制累积依赖和不超过 \(\tau\)，即有 \(\text{TV}(P_\theta(Y_S \mid X, Y_U), Q_\theta(Y_S \mid X, Y_U)) \le \tau\)。借助最大耦合理论，这严格保证了并行采样与理想联合采样的不匹配概率至多为 \(\tau\)。
损失函数 / 训练策略¶
依赖预测器的训练采用解耦的两阶段管线： 1. TV 缓存生成阶段：在 Tulu 3 SFT Mixture 数据集上抽取 100K 个响应，施加随机掩码比例 \(t \sim \text{Uniform}(0, 1)\)。针对掩码位置 \(m_j\) 采样具体 token 实现 \(y_{m_j}\)，分别前向运行骨干网络获取揭示前后的真实概率分布并计算经验全变差 \(D_{i,j}(y_{m_j})\)，共生成约 500K 组缓存样本（在 8 卡 H200 上耗时 3.5 小时）。 2. 预测器优化阶段：冻结骨干网络，仅以 float32 精度优化轻量预测头。优化目标为非对角位置的均方误差损失： $\(\mathcal{L} = \frac{1}{\sum_{b=1}^B (|M^{(b)}|^2 - |M^{(b)}|)} \sum_{b=1}^B \sum_{i \ne j} \left( \hat{D}_{i,j}^{(b)} - D_{i,j}^{(b)}(y_{m_j}^{(b)}) \right)^2\)$ 由于真实实现 \(y_{m_j}\) 采样自真实模型条件分布，MSE 的理论最优解自然收敛于期望值 \(D_{i,j}\)。训练采用 AdamW 优化器，学习率 \(10^{-5}\)，权重衰减 0.01，余弦退火配合 5% 热身，共微调 5 个 epoch（约 95 分钟）。
实验关键数据¶
主实验¶
在 Dream-7B 骨干模型上针对 MMLU-Pro、GSM8K、HumanEval 和 MBPP 进行端到端评测。所有方法统一应用 EOS 提早终止优化，基准延迟相对单步单 token 的 Entropy 基线进行归一化（平均准确率 Avg 遵从原文标准，不包含调参集 GSM8K）：
| 方法 | MMLU-Pro (%) | GSM8K (%) | HumanEval (%) | MBPP (%) | Avg (%) | 相对加速比 |
|---|---|---|---|---|---|---|
| Entropy (1 tok) | 43.9 ± 0.4 (1.0×) | 81.2 ± 1.1 (1.0×) | 53.7 ± 3.9 (1.0×) | 58.0 ± 2.2 (1.0×) | 51.8 | 1.0× |
| Entropy (2 tok) | 39.7 ± 0.4 (2.0×) | 76.7 ± 1.2 (2.2×) | 31.1 ± 3.6 (1.1×) | 43.8 ± 2.2 (1.3×) | 38.2 | 1.5× |
| Top-1 (1 tok) | 45.5 ± 0.4 (1.0×) | 81.2 ± 1.1 (1.0×) | 53.7 ± 3.9 (1.0×) | 58.0 ± 2.2 (1.0×) | 52.4 | 1.0× |
| Top-1 (2 tok) | 43.5 ± 0.4 (2.0×) | 76.7 ± 1.2 (2.1×) | 31.1 ± 3.6 (1.1×) | 43.8 ± 2.2 (1.4×) | 39.5 | 1.5× |
| Token Order (1 tok) | 45.7 ± 0.4 (1.3×) | 80.7 ± 1.1 (1.1×) | 55.5 ± 3.9 (1.3×) | 57.8 ± 2.2 (1.5×) | 53.0 | 1.3× |
| Token Order (2 tok) | 43.6 ± 0.4 (2.8×) | 73.8 ± 1.2 (2.1×) | 24.4 ± 3.4 (1.1×) | 40.8 ± 2.2 (1.5×) | 36.2 | 1.8× |
| KLASS (Kim et al., 2025) | 45.7 ± 0.4 (1.3×) | 82.0 ± 1.1 (1.1×) | 55.5 ± 3.9 (1.1×) | 58.2 ± 2.2 (1.6×) | 53.1 | 1.3× |
| DEMASK (本文) | 47.5 ± 0.4 (1.7×) | 82.8 ± 1.0 (1.8×) | 54.9 ± 3.9 (1.7×) | 57.6 ± 2.2 (2.2×) | 53.3 | 1.9× |
消融实验¶
关于是否引入 EOS 优化对解码时间（分钟）和精度的消融分析（未开启 EOS 优化时，模型会浪费前向步数反复填补序列尾部已退化的掩码位）：
| 评测基准 | DEMASK 准确率 (%) | DEMASK 耗时 (min) | KLASS 准确率 (%) | KLASS 耗时 (min) | 说明 |
|---|---|---|---|---|---|
| MMLU-Pro | 47.6 | 89.6 | 45.7 | 126.7 | 未优化 EOS 下 DEMASK 速度与准确率均全面超越 KLASS |
| GSM8K | 83.7 | 10.1 | 82.6 | 14.3 | DEMASK 维持更高准确率且耗时缩短 29.4% |
| HumanEval | 58.5 | 5.8 | 54.3 | 4.1 | KLASS 因批量无脑填充尾部空白 EOS 产生虚假加速 |
| MBPP | 56.6 | 19.3 | 57.0 | 6.9 | 代码基准最大长度预设为 1024，未优化 EOS 时 KLASS 虚高明显 |
关键发现¶
- 并行朴素扩展的灾难性退化：当无视依赖强行每步解掩 2 个 token 时，Entropy 和 Top-1 策略在 HumanEval 上准确率从 53.7% 断崖式跌落至 31.1%，而 DEMASK 能够在获得 1.7× 加速的同时维持 54.9% 的准确率。
- 次可加性假设的经验成立：在 Tulu 3 SFT 数据集上抽样的 5,000 个测试样本中，经验松弛度 \(\text{RHS} - \text{LHS}\) 统计显示违背次可加性的比例严格低于 6%，且随着子集规模 \(|S|\) 增大，平均松弛度显著上升，证实了理论界限在真实场景下的可靠性。
- 帕累托前沿全面占优：在 GSM8K 超参数网格搜索中，DEMASK 在相同扩散步数预算下准确率全面超越 KLASS，处于帕累托最优前沿。
亮点与洞察¶
- 将无向并行转化为有向依赖图剪枝：将原本棘手的多变量联合分布近似误差，转化为单次前向可解的成对全变差估计，用最小的算力代价化解了并行扩散的核心痛点。
- 权重合并消减推理 FLOPs：训练时保持双矩阵乘法分解以优化梯度动态，推理时将其融合成单一矩阵 \(\mathbf{W} = \mathbf{W}_Q \mathbf{W}_K^\top\)，避免了多次张量交互对解码延迟的侵蚀。
- 理论推导与最大耦合操作意义自洽：通过 Telescoping Sum 分布插值技巧，严密将贪心算法的累加停止条件与总变差上界挂钩，摆脱了单纯启发式调参的被动局面。
局限与展望¶
- 理论与预测误差的间隙：虽然定理基于精确成对矩阵证明了 TV 界限，但实际预测值 \(\hat{\mathbf{D}}\) 与真实值之间的泛化残差对理论界的扰动尚未建立严密泛化误差界。
- 单一骨干架构验证：论文目前仅在 Dream-7B 上完成了完整实证，尚未迁移验证至从头预训练的双向扩散大模型（如 LLaDA 8B 或 100B）。
- 静态超参数限制：当前解码过程中依赖阈值 \(\tau\) 和置信度阈值 \(\gamma\) 在整个去噪过程中保持恒定，未来若引入随时间步衰减的动态自适应阈值，或与重掩码机制协同，有望进一步压榨解码潜能。
相关工作与启发¶
- vs KLASS (Kim et al., 2025)：KLASS 依赖相继时间步之间的分布稳定性（KL 散度）做单 token 判定，本质仍是独立选择；DEMASK 则显式建模跨掩码 token 的成对条件影响，杜绝了相关性极强的 token 扎堆采样的冲突。
- vs PUNT (Azangulov et al., 2025)：PUNT 同样试图寻找条件独立集合，但需要分治迭代执行 \(\mathcal{O}(\log |M|)\) 次前向传播探查，极大稀释了并行加速红利；DEMASK 仅需单次前向即可获得整矩阵估计。
- vs APD (Israel et al., 2025)：APD 借助额外的小型自回归模型校准联合概率乘积；DEMASK 无需外挂异构语言模型，仅利用 26M 参数的微型预测头在扩散隐空间内即可自闭环。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次提出单次前向预测全变差依赖矩阵并结合贪心选择证明总变差距离上界，视角独到精巧。
- 实验充分度: ⭐⭐⭐⭐☆ 覆盖推理、数学、代码四大主流基准，提供了详尽的超参帕累托分析与次可加性实证，稍欠缺更多 8B/100B 骨干的横向迁移。
- 写作质量: ⭐⭐⭐⭐⭐ 数学证明严谨完整，算法描述与图示逻辑自洽，实验消融切中要害。
- 价值: ⭐⭐⭐⭐⭐ 为离散扩散语言模型从学术探索走向高吞吐工业化实用落地提供了极其关键的无损加速思路。