TACG：面向扩散语言模型解码的轨迹感知提交门控

# TACG：基于轨迹感知的承诺门控扩散语言模型解码
来源：https://arxiv.org/html/2607.03236
程成王¹\* 张天洛²\* 李文昊¹ 郭建元² 徐畅¹
¹悉尼大学 ²香港城市大学
[email protected]@my.cityu.edu.hk [email protected]@[email protected]

###### 摘要

扩散语言模型通过迭代去噪掩码位置来生成文本，这会暴露出一系列预测分布，而非单一的瞬时信念。大多数现有的解码器忽略了这个轨迹，仅从当前快照中确定性地承诺令牌，从而将置信度与“承诺就绪性”混为一谈：不完整上下文下的瞬时 top-1 峰值可能被锁定，而具有一致跨步骤支持的候选令牌却被延迟。我们提出了**轨迹感知承诺门控（TACG）**，一种无需训练的门控级解码器，它将令牌身份锚定到基本后验，并仅使用轨迹感知信号来决定当前提案是否准备好被承诺。TACG 结合了**时间隐式对数引导（TILG）**，该模块维护过去对数的一个指数移动平均作为自参考，并在自然参数空间中将当前对数与该参考进行对比，以及一个**历史门（HG）**，该门在承诺前强制执行短期提案持久性。再加上一个上限的额外提升预算，这些组件形成了一个稳定性约束的承诺规则，无需辅助网络或额外的前向传播。我们在代码（HumanEval, MBPP）和数学（GSM8K, MATH500）基准上对 LLaDA、Dream 和 LLaDA2-Mini 进行了评估；TACG 通常在提高或保持准确率的同时，减少了去噪步数并增加了每次前向传播的令牌数（TPF）。代码已公开在 https://github.com/Clarence-CV/TACG-DLLM。

¹ 共同第一作者。

## 1 引言

扩散语言模型通过反转掩码过程[2](https://arxiv.org/html/2607.03236#bib.bib2), [30](https://arxiv.org/html/2607.03236#bib.bib30), [21](https://arxiv.org/html/2607.03236#bib.bib21), [23](https://arxiv.org/html/2607.03236#bib.bib23)生成文本。从许多位置被掩码的序列开始，模型反复预测掩码位置，并逐步将其中一部分承诺为具体令牌。与自回归模型一次扩展一个令牌不同，DLLM 的单次前向传播会同时为多个位置生成预测分布。因此，解码不仅是一个令牌预测问题，也是一个决定哪些预测足够可靠以进入部分生成序列的决策问题。

随着去噪过程的进行，某个位置的对数会随周围上下文变化而变化：排名第一（top-1）的候选令牌、其置信度以及分布形态都可能发生改变。因此，每个掩码位置都承载着一个时间信念轨迹，而不仅仅是一个瞬时预测。

大多数当前的 DLLM 解码器对此轨迹的利用有限。基于置信度的重新掩码在每个步骤承诺置信度最高的位置。部分继承自掩码生成建模[4](https://arxiv.org/html/2607.03236#bib.bib4)的基于调度的解码器，会揭示预定数量的位置，而不考虑当前模型状态。稳定性感知解码器，如 KLASS[15](https://arxiv.org/html/2607.03236#bib.bib15)，承诺其预测分布看起来局部稳定的位置。这些规则大致基于快照：在步骤 t 的决策由当前分布 p_t 决定，可能辅以局部稳定性统计。模型产生了一个轨迹，但解码器仅消费其最新状态。

这种快照视角对于无需训练的 DLLM 解码来说是不够的，其目标是在不改变模型参数的情况下改善质量-效率的权衡。解码器应避免不可靠的早期承诺，因为一个错误解掩的令牌会成为后续去噪步骤的条件上下文的一部分。同时，它应尽早揭示就绪的令牌，因为延迟它们会增加去噪步数并降低并行性。简单地调整置信度阈值无法解决这种紧张关系：更严格的阈值提高了谨慎性，但延迟了就绪的令牌；而更宽松的阈值加速了解码，但接纳了不稳定的候选者。

这自然引出一个问题：一个无需训练的解码器能否从去噪过程中已有的信号中估计承诺就绪性，而不是通过重新训练模型或添加前瞻性？我们发现历史轨迹提供了这样的信号。如图1（https://arxiv.org/html/2607.03236#S1.F1）所示，我们在 GSM8K 上的诊断分析揭示了两种现象，详见第2节（https://arxiv.org/html/2607.03236#S2）。在可比置信度下，历史门支持提高了令牌级匹配，表明存在超出快照置信度的可靠性线索。历史对数支持在较低置信度区域尤其有用，表明在刚性置信度阈值揭示它们之前，历史对数可以识别出准备就绪的承诺令牌。这些信号共同支持了可靠的承诺和早期的令牌提升。

参考图注：图1：GSM8K 上的诊断历史信号。更高的历史门和历史对数支持改善了令牌级匹配，表明它们在可靠承诺和早期令牌提升中发挥作用。
基于这一观察，我们将快照置信度常常合并的两个决策分离开来：应该写入哪个令牌，以及该令牌何时应该离开掩码状态。瞬时置信度衡量的是当前后验峰的高度，但它并不表明该峰值是瞬变的、停滞的，还是得到去噪轨迹支持的。因此，我们将令牌身份锚定到基本后验，并仅使用历史轨迹来估计承诺就绪性。具体来说，基本模型提出 x̂_i = argmax_v p_t(i,v)，置信度 c_i = p_t(i, x̂_i)，TACG 决定这个提案是否现在应该被承诺。

我们引入了**轨迹感知承诺门控（TACG）**作为一种门控级解码框架。TACG 结合了两个互补的轨迹感知组件。其时间支持分支，**时间隐式对数引导（TILG）**，维护过去对数的一个指数移动平均 z̄_{t-1} 作为隐式自参考，并使用当前信念与该参考之间的对比来计算基本提案的时间支持。其稳定性分支，**历史门（HG）**，要求在承诺前有短期提案持久性，并为已经饱和的位置设置了置信度逃逸。为简洁起见，定义 TILG 辅助读出

q_t(v) ∝ p_t(v) (p_t(v) / p_t^ref(v))^w = p_t(v)^{1+w} / p_t^ref(v)^w. (1)

在操作上，TILG 分支仅查询基本提案 x̂_i 处的该读出值：如果该提案的支持度相对于历史参考有所增加，其就绪性分数上升；如果当前后验峰值是瞬变的或正在减弱，时间支持仍然很小。这种对比是在单个模型内的去噪时间上进行的，而参考是模型自身的缓慢信念，来自解码过程中已产生的对数。因此，TILG 不需要辅助网络，也不需要额外的模型前向传播。

为了获得实用的解码器，TACG 将这一时间信号与一个稳定性约束的承诺规则相结合。HG 通过要求在承诺前有一个短的持久窗口来抑制瞬变候选者，同时一个置信度下限排除了明显未准备好的候选者。在稳定基本接受集之外，TACG 提升一个有界数量的额外位置，这些位置具有最大的就绪性分数 s_i = c_i + λ σ_i，其中 c_i 是提案置信度，σ_i 是提案令牌的时间支持。上限 K_extra 作为一个明确的加速预算：时间证据可以在步内移动揭示边界，但仅在受控的额外承诺数量内。

我们的贡献如下。首先，我们识别出承诺时机作为 DLLMs 中一个独立的解码器决策，并解释了为什么瞬时置信度是承诺就绪性的不完整代理。其次，我们引入了 TACG，一个将令牌身份与承诺时机分离的门控级框架，并结合了 TILG 时间支持与 HG 持久性约束。第三，我们将这些组件与一个稳定性约束的加速规则耦合，该规则具有上限的、一致性条件化的额外提升。第四，我们在代码[5](https://arxiv.org/html/2607.03236#bib.bib5), [1](https://arxiv.org/html/2607.03236#bib.bib1)和数学[7](https://arxiv.org/html/2607.03236#bib.bib7), [11](https://arxiv.org/html/2607.03236#bib.bib11)基准上对 LLaDA[23](https://arxiv.org/html/2607.03236#bib.bib23), Dream[41](https://arxiv.org/html/2607.03236#bib.bib41)和 LLaDA2-Mini[3](https://arxiv.org/html/2607.03236#bib.bib3)评估了 TACG，展示了在基于置信度的门控下的一致准确率提升和效率提升。

## 2 预备知识

掩码扩散语言模型（MDLMs）。MDLMs 将文本生成视为一个序列上的离散去噪问题。令 x_0 = (x_0^1, ..., x_0^L) ∈ V^L 为一个长度为 L 的干净令牌序列，词汇表为 V。在前向加噪过程中，x_0 中的令牌逐渐被损坏成一个特殊的掩码符号 [MASK]。给定一个在噪声水平 t ∈ [0,1] 下部分损坏的序列 x_t，模型学习预测掩码位置处的原始干净令牌。一个常用的 MDLM 训练目标可写为

L_MDLM(θ) = -E_{x_0, t, x_t} [ (1/t) Σ_{i=1}^L 1[x_t^i = [MASK]] log p_θ(x_0^i | x_t) ]. (2)

在推理时，生成通常从完全掩码的序列开始。模型然后并行地反复预测掩码位置的令牌，可选地再次掩码不确定的预测以允许进一步细化。

诊断性历史信号。我们在 GSM8K 诊断子集上分析两个历史信号，以研究承诺就绪性是否可以在快照置信度之外更好地估计。门级持久性衡量提案是否在步骤间保持稳定，而对数级强化衡量提案是否从历史对数轨迹获得支持。

• **历史门支持**。我们首先检查历史证据的一种离散形式：同一个令牌提案是否在相邻去噪步骤间持续存在。对于每个掩码位置 i，我们将**历史门支持**定义为当前提案保持不变连续步数：

ℓ_{t,i} = { ℓ_{t-1,i} + 1,  x̂_{t,i} = x̂_{t-1,i}; 1, x̂_{t,i} ≠ x̂_{t-1,i} } (3)

较大的 ℓ_{t,i} 表示提案已在去噪轨迹上持续了更长时间。图1（左）显示，在匹配的置信度区域内，具有更高历史门支持的候选者实现了更高的令牌级匹配。这表明提案持久性提供了超出快照置信度的可靠性线索。换句话说，即使两个候选者具有相似的置信度，在最近去噪步骤中保持稳定的那个更值得承诺。这激发了使用历史门支持作为门级信号来提高承诺质量。

• **历史对数支持**。虽然门级持久性衡量提案是否稳定，但提高解码效率也需要识别在达到刚性置信度阈值之前可以承诺的候选者。简单地降低置信度阈值可以更早揭示更多令牌，但它平等对待所有较低置信度的候选者，并可能引入不稳定的承诺。因此，我们检查去噪轨迹中的一个连续对数级信号。

令 z̄_{t-1}(i,·) 为位置 i 处先前对数的 EMA 参考。给定当前对数 z_t(i,·)，我们形成一个历史引导的读出

q_t(i,·) = softmax( z_t(i,·) + λ[z_t(i,·) - z̄_{t-1}(i,·)] ), (4)

其中 λ 控制历史信号的强度。对于当前提案 x̂_{t,i}，我们通过该提案被分配的增益来衡量其历史对数支持：

g_{t,i} = q_t(i, x̂_{t,i}) - p_t^ref(i, x̂_{t,i}),   p_t^ref(i,·) = softmax(z̄_{t-1}(i,·)). (5)

这个信号是候选条件化的：它评分当前基本提案是否被历史对数轨迹强化，而不是用一个新令牌替换它。图1（右）显示，具有更高历史对数支持的候选者实现了更高的令牌级匹配，尤其是在较低置信度区域。这表明一些低于刚性置信度阈值的候选者已经得到了去噪轨迹的支持。因此，历史对数支持为早期令牌提升提供了一个对数级信号：它可以有选择地扩展承诺边界，而不是均匀地放松置信度阈值。

## 3 相关工作

扩散语言模型。扩散模型[13](https://arxiv.org/html/2607.03236#bib.bib13),[32](https://arxiv.org/html/2607.03236#bib.bib32)已成为视觉生成中的主流[28](https://arxiv.org/html/2607.03236#bib.bib28),[25](https://arxiv.org/html/2607.03236#bib.bib25),[29](https://arxiv.org/html/2607.03236#bib.bib29),[46](https://arxiv.org/html/2607.03236#bib.bib46)，并且最近的工作探索了它们在文本生成中的应用。在现有范式中，掩码扩散语言模型（MDLMs）[31](https://arxiv.org/html/2607.03236#bib.bib31),[2](https://arxiv.org/html/2607.03236#bib.bib2),[30](https://arxiv.org/html/2607.03236#bib.bib30),[50](https://arxiv.org/html/2607.03236#bib.bib50),[21](https://arxiv.org/html/2607.03236#bib.bib21)通过在离散空间中通过掩码令牌预测建模语言，已成为 AR-LLM 的一个有前途的替代方案。基于此公式，LLaDA[23](https://arxiv.org/html/2607.03236#bib.bib23)和 Dream[41](https://arxiv.org/html/2607.03236#bib.bib41)通过大规模预训练将 MDLM 扩展到十亿参数规模，展示了它们的实际潜力。LLaDA-2.0[3](https://arxiv.org/html/2607.03236#bib.bib3)和 LLaDA-MoE[54](https://arxiv.org/html/2607.03236#bib.bib54)进一步表明，MDLM 可以通过混合专家架构有效地扩展。除了这些进展，dLLMs 也正吸引越来越多的关注。