Mean-Field Parallel Decoding for Discrete Diffusion Language Models¶

会议: NeurIPS 2026
arXiv: 2606.15805
论文: https://arxiv.org/abs/2606.15805
代码: https://github.com/tamim-zoabi/Mean-Field-Parallel-Decoding
领域: LLM效率
关键词: 离散扩散语言模型, 均场变分推断, 并行解码, 联合不一致, 杰森-香农散度

一句话总结¶

针对离散扩散语言模型（dLLM）在并行解掩码时因边缘独立贪心更新造成的联合不一致与词元冲突问题，本文提出一种免训练的单次前向均场并行解码框架，利用预测分布间的 JSD 散度显式构建成对交互抑制场，通过变分不动点迭代自适应协同并行提交决策。

研究背景与动机¶

自回归大语言模型（ARM）凭借从左至右的逐词元生成在各类任务中取得了卓越表现，但严格串行的生成范式要求每个 token 都经历一次完整的模型前向传播，导致推理延迟高居不下、难以满足实时交互需求。为了打破这一计算瓶颈，离散扩散语言模型（dLLM，如 LLaDA、Dream 等）近年来迅速崛起。dLLM 将文本生成形式化为反向离散去噪过程，允许解码器根据预定策略在单步去噪中同时解除多个掩码位置，为实现高吞吐、低延迟的非自回归与半自回归解码开辟了可行路径。

然而，激进的并行解码在实践中往往伴随着严重的「联合不一致性」（Joint Inconsistency）或「联合失配」（Joint Mismatch）。绝大多数现有 dLLM 采用完全因子分解的反向去噪策略，在单步去噪中独立预测每个被掩码位置的词元分类分布。这种独立性假设忽略了自然语言中无处不在的句法依存与语义协调约束；当解码器仅仅依据各位置的边缘局部置信度（如预测熵或 Top-1/Top-2 margin）贪心地选出多个位置同时提交时，那些孤立看来高度可信的词元组合在一起却往往相互冲突，极易引发连续重复同一词元、句法错乱或语义断裂。

现有的缓解方案各存局限：DAWN 或 DAPD 等方法依赖注意力图启发式构建或额外的跨词元依赖估计；DEMASK 引入了额外的轻量预测器网络进行两两影响估计；块扩散（Block Diffusion）则受限于半自回归块级约束，且无法解决块内并行的依赖冲突。这些方法要么依赖辅助模型、引入额外的重训练负担，要么需要多次额外计算。本文的核心切入角度在于：利用模型自身在当前单次前向中产出的多位置预测分布重叠度，直接定义词元间的互斥竞争关系，并借由经典的变分均场理论在单步内极速完成协同决策。核心 idea：将离散扩散的并行提交子集选择形式化为成对马尔可夫随机场（MRF）上的结构化推断问题，以局部对数边际为一阶势能、以预测分布间的规范化 JSD 散度为反铁磁二阶抑制势能，通过极轻量的均场定点迭代（$R=2$）在单次前向中自适应压制冲突提交。

方法详解¶

整体框架¶

在离散扩散模型的去噪步 $t \in \{T, \dots, 1\}$ 中，记当前未解掩的位置集合为 $\mathcal{U}_t$，目标是从中筛选出一个互不冲突且高置信度的并行提交子集 $\mathcal{C}_t \subseteq \mathcal{U}_t$（即令指标向量 $S \in \{0, 1\}^m$，其中 $m = |\mathcal{U}_t|$）。整个均场并行解码流水线分为四个阶段：首先由模型单次前向输出各掩码位置的类别分布 $\pi_{t,i}$，并计算 Top-2 对数边际作为局部置信度；随后基于所有掩码位置分布间的 Jensen-Shannon 散度（JSD）构建成对竞争抑制矩阵 $D$；接着运行 $R$ 轮均场变分不动点迭代协同更新提交强度向量 $q$；最后按提交阈值 $\tau$ 截断强度，提取出最终的并行解掩集合 $\mathcal{C}_t$。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["单步前向传播输出<br/>获取掩码集合 Ut 的 logits 与分布 πt,i"] --> B["1. 局部置信度估计<br/>计算 Top-2 对数边际 ci 作为一阶势能"]
    A --> C["2. 成对预测重叠构建<br/>基于分布间 JSD 计算反铁磁抑制矩阵 D"]
    B --> D["3. 变分均场迭代松弛<br/>固定点迭代 q(r+1) = σ(c - Dq(r))，R=2"]
    C --> D
    D --> E["4. 强度阈值截断筛选<br/>保留 qi ≥ τ 的非冲突子集 Ct 并行解掩"]

关键设计¶

1. 局部置信度一阶势能构建：以对数边际量化单位置确定性

针对传统仅看单一概率最大值容易忽略近邻候选竞争的问题，本文将指示变量 $s_i \in \{0, 1\}$ 的一阶势能设计为线性项 $\phi_i(s_i) = c_i s_i$。局部置信度 $c_i$ 严格定义为预测分布 $\pi_{t,i}$ 中概率最高词元 $v_i^{(1)}$ 与次高词元 $v_i^{(2)}$ 之间的对数差值（Log-margin）： $$c_i = \log \pi_{t,i}(v_i^{(1)}) - \log \pi_{t,i}(v_i^{(2)})$$ 该设计使得当某一位置具有压倒性的预测优势时，$c_i$ 呈现显著正值，为该位置提供强烈的积极提交倾向；而对于预测模糊、模态分散的不确定位置，$c_i$ 接近于 0 甚至很小，自然压低其初始提交权重，从而在底层保证了高可信候选的优先解掩。

2. 基于 JSD 散度的成对反铁磁抑制势能：量化词元更新竞争

针对并行解码中最致命的“相邻或关联位置同时预测相似词元造成冗余或冲突”的痛点，本文在成对势能中引入二阶惩罚项 $\phi_{ij}(s_i, s_j) = -D_{ij} s_i s_j$。为了以极低开销在单次前向内捕获词元冲突，设计利用对称且有界的 Jensen-Shannon 散度（JSD）来度量两位置预测分布 $\pi_{t,i}$ 与 $\pi_{t,j}$ 的重叠程度： $$\tilde{D}_{ij} = 1 - \frac{\text{JSD}(\pi_{t,i}, \pi_{t,j})}{\ln 2}, \quad \text{其中 } \text{JSD}(p, q) = \frac{1}{2}\text{KL}\left(p \parallel \frac{p+q}{2}\right) + \frac{1}{2}\text{KL}\left(q \parallel \frac{p+q}{2}\right)$$ 随后对活跃集合进行自归一化 $D_{ij} = \tilde{D}_{ij} / \max_{a,b} \tilde{D}_{ab}$（对角线强制置零 $D_{ii} = 0$ 以免自抑制）。这种物理学上的“反铁磁”（Antiferromagnetic）耦合机制意味着：若两个被掩码位置倾向于预测重叠度极高或相同模式的词元（$D_{ij}$ 接近 1），同时将它们设为 1 将招致巨大的能量惩罚；反之，若两位置语义互补、分布正交，则二阶抑制近乎消失。这一设计完全无需训练辅助网络或提取注意力图，纯粹由输出空间的预测一致性自然驱动。

3. 变分均场不动点迭代松弛：极低开销的多候选全局协同

针对在任意成对矩阵 $D$ 下求解精确最大后验（MAP）配置属于 NP-hard 组合优化的难题，本文采用完全因子分解的变分推断分布 $q(S) = \prod_{i=1}^m q_i^{s_i} (1 - q_i)^{1 - s_i}$ 对真实后验进行松弛逼近，其中 $q_i \in [0, 1]$ 表示位置 $i$ 的提交强度。通过最大化证据下界（ELBO）： $$\mathcal{L}(q) = \sum_{i=1}^m c_i q_i - \frac{1}{2} \sum_{i \neq j} D_{ij} q_i q_j + \sum_{i=1}^m \left[-q_i \log q_i - (1 - q_i) \log(1 - q_i)\right]$$ 对每个分量求偏导并令 $\frac{\partial \mathcal{L}}{\partial q_i} = 0$，直接导出形式极其优雅的不动点迭代方程： $$q_i^{(r+1)} = \sigma\left(c_i - \sum_{j \neq i} D_{ij} q_j^{(r)}\right)$$ 其中 $\sigma(\cdot)$ 为标准 Sigmoid 函数。以单位置置信度 $q^{(0)} = \sigma(c)$ 为初值，各位置根据周围已激活节点的竞争抑制动态衰减自身强度。实证表明，该迭代仅需 $R=2$ 轮即可收敛至稳定不动点。最终只需通过设定的标量阈值 $\tau$ 进行切分即可获得协同解掩集合 $\mathcal{C}_t = \{i \in \mathcal{U}_t: q_i^{(R)} \ge \tau\}$。

一个完整示例¶

假设当前存在 3 个掩码位置（$m=3$），模型单次前向输出后计算得到的对数边际为 $c = [4.0, 3.8, 1.0]$。此时位置 1 和位置 2 局部置信度均极高，且模型对两者预测的高概率词元几乎完全相同（例如均强烈偏向预测同一重复名词），导致归一化分布重叠 $D_{12} = D_{21} = 0.95$；而位置 3 与其他两位置正交，矩阵其它互斥项接近 0。

初始提交倾向：若采用传统边缘贪心解码，位置 1 和位置 2 的局部概率都极大，将被同时提交，导致连续重复错误。在本文方法中，初值 $q^{(0)} = \sigma(c) \approx [0.98, 0.98, 0.73]$。
第一轮均场迭代：
- 位置 1 受位置 2 强抑制：有效输入为 $c_1 - D_{12} q_2^{(0)} = 4.0 - 0.95 \times 0.98 \approx 3.07$；
- 位置 2 受位置 1 强抑制：有效输入为 $c_2 - D_{21} q_1^{(0)} = 3.8 - 0.95 \times 0.98 \approx 2.87$。
第二轮均场迭代（$R=2$）：微弱的初始优势（$c_1 > c_2$）在二阶反铁磁抑制中被非线性放大。位置 1 保持较高的有效激活，而位置 2 的净输入持续被大幅抵消；若设定提交阈值 $\tau = 0.85$，经过 2 轮更新后，位置 1 顺利保留在提交集 $\mathcal{C}_t$ 中，而位置 2 的强度被压制在 $\tau$ 之下推迟至后续去噪步更新，成功化解了并发生词冲突。

实验关键数据¶

主实验¶

在四项主流推理与代码基准（GSM8K、MATH、HumanEval、MBPP）上，采用三大代表性 dLLM 骨干（LLaDA-8B-Instruct、LLaDA-1.5、Dream-v0-Instruct-7B）对本文方法与主流解码策略进行了全方位对比。评测指标包含任务准确率（Acc.）、每秒生成词元数（TPS）以及相对 Entropy 标杆的加速比（Speedup）。

骨干模型与解码方法	GSM8K Acc. / TPS / 加速比	MATH Acc. / TPS / 加速比	HumanEval Acc. / TPS / 加速比	MBPP Acc. / TPS / 加速比
LLaDA-8B-Instruct
Entropy	79.15 / 15.87 / 1.00×	33.36 / 23.59 / 1.00×	40.24 / 38.02 / 1.00×	29.40 / 15.72 / 1.00×
KLASS	75.05 / 36.23 / 2.27×	31.80 / 42.20 / 1.79×	39.63 / 70.30 / 1.83×	27.20 / 35.60 / 2.28×
LocalLeap	77.71 / 68.96 / 4.32×	32.96 / 77.87 / 3.30×	40.24 / 152.62 / 4.04×	30.60 / 72.64 / 4.60×
DAWN	77.78 / 69.76 / 4.38×	32.34 / 80.40 / 3.41×	40.24 / 152.46 / 4.05×	28.00 / 74.46 / 4.87×
Ours	76.34 / 78.75 / 4.93×	31.10 / 96.37 / 4.07×	40.70 / 180.81 / 4.88×	28.00 / 81.11 / 6.10×
LLaDA-1.5
Entropy	81.80 / 14.90 / 1.00×	33.62 / 20.91 / 1.00×	44.51 / 13.04 / 1.00×	39.20 / 5.77 / 1.00×
KLASS	79.37 / 35.22 / 2.37×	32.50 / 38.52 / 1.87×	43.90 / 18.94 / 1.48×	29.80 / 18.25 / 4.08×
LocalLeap	80.36 / 65.69 / 4.41×	33.02 / 68.87 / 3.29×	42.07 / 41.91 / 3.23×	38.00 / 43.50 / 7.57×
DAWN	79.60 / 66.74 / 4.47×	32.16 / 70.90 / 3.39×	42.68 / 41.50 / 3.20×	37.40 / 45.36 / 7.90×
Ours	80.74 / 75.18 / 4.96×	32.46 / 86.71 / 4.10×	42.24 / 54.91 / 4.32×	38.60 / 50.95 / 8.62×
Dream-v0-Instruct-7B
Entropy	75.51 / 11.15 / 1.00×	38.30 / 28.18 / 1.00×	53.04 / 32.05 / 1.00×	54.20 / 13.40 / 1.00×
KLASS	73.54 / 19.61 / 1.82×	36.64 / 29.26 / 1.70×	53.65 / 38.63 / 1.83×	56.60 / 29.00 / 1.30×
LocalLeap	73.61 / 49.80 / 4.49×	38.38 / 68.62 / 2.51×	56.09 / 84.42 / 3.05×	55.00 / 66.86 / 5.39×
DAWN	73.31 / 50.15 / 4.53×	38.42 / 69.77 / 2.57×	54.87 / 89.56 / 3.26×	55.60 / 69.50 / 5.62×
Ours	74.90 / 54.75 / 4.91×	40.44 / 74.38 / 3.61×	—	54.40 / 75.68 / 5.80×

消融实验¶

基于 Dream-v0-Instruct-7B 骨干在 GSM8K 数据集上（默认配置 $\tau = 0.90, R = 2, \text{Block size} = 32$）对算法核心组件进行系统剥离消融，统计准确率、TPS 以及平均去噪函数评估步数（NFE）：

配置变体	准确率（Acc. ↑）	吞吐量（TPS ↑）	去噪步数（NFE ↓）	变体设计说明
Baseline (Entropy)	75.51	11.15	256.0	标准按熵保守串行/半并行去噪基线
No pairwise interaction	72.40	56.20	48.7	移除成对势能项（$D = 0$），纯靠局部边际 $c_i$ 阈值筛选
Uniform interaction matrix	71.95	51.92	52.6	将 JSD 矩阵替换为非对角线全为常数的均匀成对抑制阵
One-shot update ($R=1$)	54.51	20.11	140.0	仅执行单步均场迭代更新，尚未收敛至协同不动点
Full Model (Ours)	73.24	49.13	54.3	完整 JSD 矩阵 + $R=2$ 轮均场不动点迭代松弛

关键发现¶

成对抑制是防止性能崩塌的核心支柱：在消融实验中，完全剔除成对抑制项（No pairwise interaction）虽然吞吐量略有提升，但准确率直接从 73.24% 掉至 72.40%；而若使用均匀抑制矩阵（Uniform），准确率进一步骤降至 71.95%。这证明由 JSD 捕获的精细几何结构与反铁磁竞争至关重要，单纯的通用惩罚反而损害了正交独立词元的安全并行。
均场迭代需达到最低收敛阶数：仅迭代一次的 $R=1$ 配置性能遭遇剧烈雪崩（准确率大幅暴跌至 54.51%，NFE 飙升至 140.0），原因在于单步更新下各位置强度剧烈震荡且未形成有效协商，导致大量位置被错误抑制或冲突误判；而设置 $R=2$ 即可实现优异的定点平衡，进一步增加到 $R=4, 6, 8$ 时准确率无明显提升反而因微量 CPU/GPU 开销导致 TPS 略有下降。
广泛的吞吐量与加速比优势：在全套评估测试中，本文方法在所有配置下均斩获了最高的 TPS，平均加速比达到 5.12×（最高在 LLaDA-1.5 的 MBPP 上达到 8.62× 加速），且在多项任务上（如 Dream-7B 在 MATH 上达 40.44% vs Baseline 38.30%，LLaDA-8B 在 HumanEval 上达 40.70% vs Baseline 40.24%）准确率持平甚至超越了保守的 Entropy 基准。

亮点与洞察¶

将扩散解掩离散决策精妙转化为 MRF 均场变分推断：跳出了传统启发式设定阈值或依赖注意力矩阵的固有框架，从统计物理与图模型能量角度将“高置信度追求”与“冲突惩罚”统一于连续变分下界，通过数学推导直接给出无参数的 Sigmoid 不动点迭代。
JSD 散度作为分布重叠的天然代理指标：无需构建注意力依赖图，亦无需像 DEMASK 那样训练辅助神经预测器，直接对当前单次前向输出的 logits 计算对称且有界的 JSD 散度，完美规避了额外模型前向与反向传播的延迟开销。
极佳的免训练即插即用性与高度工程友好度：整个方法仅包含数十行张量运算，完全运行在 GPU 单步前向后的后处理阶段；$O(m^2 |\mathcal{V}|)$ 的 JSD 计算在块扩散尺寸下高度并行化，单次迭代 $R=2$ 的开销相对于 Transformer 主干而言近乎可以忽略。

局限与展望¶

计算复杂度随掩码候选集合呈二次增长：成对 JSD 矩阵构建需要 $O(m^2 |\mathcal{V}|)$ 的时间与显存复杂度。在常规块级解码（Block size 20~32）下非常高效，但若在超长序列（如 $m \ge 256$）或超大词表模型中进行无块限制的全局并行解码，JSD 计算开销将显著增加，未来需探索基于 Top-k 截断或低秩投影的快速近邻 JSD 近似。
成对相互作用无法捕获更高阶的语义组合依赖：MRF 仅建模了两两位置间的重叠惩罚，无法直接表达三元组及以上更复杂的长程句法闭环或逻辑依存关系。
超参数敏感性依然存在：提交阈值 $\tau$ 与块尺寸大小直接主导了速度与精度的 Pareto 前沿，在不同架构或特定下沉任务上仍需少量子集样本进行调优以匹配最优工作点。

评分¶

新颖性: ⭐⭐⭐⭐☆ 首次将均场变分理论与 JSD 分布距离引入离散扩散解码，构建了典雅的二阶冲突抑制模型。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 大主流 dLLM 骨干、4 项经典数学与代码推理基准，并附带详尽的消融与参数敏感度分析。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰严密，从联合不一致的问题痛点到变分下界推导与实验论证一气呵成。
价值: ⭐⭐⭐⭐⭐ 免训练、单次前向、即插即用，大幅刷新了离散扩散语言模型的推理速度与实用化落地水平。