路径至关重要：为扩散语言模型学习令牌提交策略

# 路径至关重要：为扩散语言模型学习令牌提交策略
来源：https://arxiv.org/html/2605.24697
Bohang Sun¹ * Max Zhu¹ † Francesco Caso¹ Jindong Gu² Junchi Yu² Philip Torr² Pietro Liò¹ Jialin Yu²  
¹剑桥大学计算机科学与技术系 ²牛津大学工程科学系

###### 摘要

扩散大语言模型通过并行精炼多个令牌位置来提升生成速度，但这种并行性引入了一个隐藏的控制问题：在每一步，哪些被提出的令牌应该被转移到部分解码序列中？我们将这一决策称为*令牌提交*。现有的冻结生成器解码器主要依赖人工设计的置信度规则或特定模块的接受过滤器。本文论证令牌提交可以作为一种可重用的轨迹状态策略来学习。我们提出**TraceLock**，一种轻量级插件式控制器，为冻结的扩散语言模型实例化该策略。由于无法获得最优提交时间，TraceLock 从未来稳定性中获取自监督信号：在解码步骤 *t*，如果位置 *i* 上的被提议令牌与完整解码轨迹结束后位置 *i* 的最终令牌一致，则被标记为稳定。该控制器对可变长度轨迹状态进行评分，并决定哪些活跃的令牌提议应被提交到部分解码序列中。一旦为给定冻结骨干模型完成训练，该控制器即可跨局部窗口宽度、生成长度和步数预算进行部署，无需重新训练或针对每个设置进行校准。在问答、数学推理和代码生成任务上的实验表明，TraceLock 在质量-步数权衡上优于启发式方法和学习型基线，尤其在跨设置部署下表现出特别稳定的行为。诊断分析显示，其决策无法简化为标量置信度，这表明冻结扩散语言模型暴露了一个超越置信度解码的可学习提交轨迹空间¹。

¹ 我们的代码已发布在 https://github.com/BobSun98/TraceLock。

## 1 引言

扩散大语言模型（D-LLMs）已成为自回归大语言模型（AR-LLMs）的一种有前景的替代方案，尤其是在掩码离散扩散设置中（Sahoo 等人，2024；Ou 等人，2025），如 LLaDA 和 Dream（Nie 等人，2025；Ye 等人，2025）。与从左到右不可逆地附加令牌的 AR-LLMs 不同，掩码 D-LLMs 迭代地精炼一个部分掩码的序列，并且可以并行更新多个令牌位置。这种并行精炼避免了自回归解码的固定生成顺序，并创造了更快生成的可能性。然而在实践中，高质量的 D-LLM 生成通常仍然需要多次去噪迭代，每次迭代都涉及对当前序列的双向注意力。因此，高效解码是有效部署 D-LLMs 的核心挑战（Wu 等人，2025b；Chen 等人，2025b；Israel 等人，2025）。

D-LLMs 的快速解码通常被描述为减少去噪步骤的数量或成本。然而，在冻结生成器设置中，生成器是固定的，主要的算法自由度在于其他地方：解码器必须决定哪些被提议的令牌应停止被修订。*这便将快速解码转化为一个令牌提交问题*。在每个去噪步骤，冻结模型为多个位置提议令牌值，而解码器必须决定哪些位置应被提交，哪些应保持可修订性。选择有效的提交策略并非易事：提交过晚限制加速效果，提交过早则可能锁定错误。图 1 说明了这种权衡。现有的冻结生成器解码器可以理解为针对特定设置校准提交的各种方式。无训练方法使用人工设计的置信度阈值、转移规则或模块调度（Wu 等人，2025b, a；Dong 等人，2025），而学习型过滤器（如 Learn2PD）则学习一个与固定模块级解码接口绑定的接受规则（Bao 等人，2026）。这自然引出一个问题：*令牌提交是否可以作为一种可重用的轨迹状态策略来学习，从而能够跨局部窗口宽度、生成长度和步数预算进行部署？*

学习这样的策略是困难的，因为最优的提交决策不可直接观测。一个完整的生成告诉我们最终的序列，但并不能告诉我们每个令牌应在何时停止被修订。因此，我们将提交学习转化为一个从完整扩散轨迹推导出的自监督轨迹预测问题。一个中间令牌提议被标记为*未来稳定*，如果它与最终完整序列中对应位置的令牌一致。这个未来稳定性目标不需要人工标注或任务级别的正确性标签。尽管它只是一个轨迹相关的信号，但它为在线决策——哪些活跃令牌可以安全提交——提供了一个密集的代理。我们提出 **TraceLock**，一种轻量级插件式控制器，从这些标签中学习可重用的提交策略。基础 D-LLM 保持冻结，TraceLock 仅决定哪些活跃生成位置应保持可修订，哪些应被锁定。TraceLock 不是设计或学习特定于设置的接受校准，而是使用冻结模型的隐藏状态、短程隐藏状态动力学以及提示/活跃/锁定上下文对可变长度轨迹状态进行评分，使用单一的共享上下文评分器。我们在问答、数学推理和代码生成任务上，在多种生成长度和局部窗口机制下评估 TraceLock。结果表明，轨迹监督的上下文提交在多个设置中改善了质量-步数权衡，优于启发式和学习型解码基线。

（参见图标题）
图 1：掩码扩散解码中令牌提交作为轨迹选择问题。激进的解码可能很快，但可能锁定错误的轨迹；而保守的解码则让令牌保持更长时间的可修订性，但成本更高。有用的操作点并非单一的全局截止点；它可能随着部分轨迹的演变在不同样本和不同步骤之间变化。

#### 贡献。
我们做出以下贡献：
**问题形式化。** 我们将高效的冻结生成器 D-LLM 解码形式化为一个令牌提交问题，其中解码器决定被提议的令牌何时应变为不可逆。
**算法。** 我们展示了这种提交策略可以从完成的扩散轨迹中使用自监督的未来稳定性标签来学习，并将其实例化为 TraceLock，一个端到端学习的提交控制器，而非人工设计或特定于模块的接受规则。
**实证证据。** 在数学推理、问答和代码生成任务上，我们展示了 TraceLock 在质量-步数权衡上优于启发式和学习型基线，在局部窗口宽度和生成长度变化下保持稳定，并且学习到的提交行为超越了标量置信度过滤。

## 2 相关工作

#### D-LLM 解码加速。
离散和掩码扩散语言模型已成为自回归语言建模日益重要的替代方案（Sahoo 等人，2024；Ou 等人，2025；Nie 等人，2025；Ye 等人，2025）。现有的加速方法在生成管道的不同层面进行干预。系统方法通过缓存或执行优化来降低单个去噪迭代的成本（Wu 等人，2025b；Jiang 等人，2025）。模型适应方法修改生成器或其训练目标，使得更快或更激进的并行解码变得可靠，例如通过微调、蒸馏、近似联合采样或学习的并行解码行为（Chen 等人，2025b；Bansal 和 Sanghavi，2025；Israel 等人，2025）。其他方法通过混合、分块或强制风格公式来改变生成过程本身（Wang 等人，2025b；Arriola 等人，2025）。这些方向与我们的工作互补：TraceLock 不训练新的扩散骨干网络，也不降低单个去噪步骤的成本，而是学习冻结生成器的推理时策略，该策略决定哪些被提议的令牌应停止被修订。

#### 冻结生成器解码加速。
我们的工作最接近于那些保持 D-LLM 冻结，并通过改变令牌接受、重新掩码或转移策略来加速解码的方法。无训练解码器（如 Fast-dLLM 及其后续变体）使用置信度阈值转移、模块调度或相关人工设计的规则来决定哪些位置应被揭示或最终确定（Wu 等人，2025b, a）。更广泛的工作研究令牌排序、置信度校准、时间建模或重新掩码行为的替代启发式方法（Li 等人，2025；Wang 等人，2025a；Kim 等人，2025；Hong 等人，2025；Dong 等人，2025）。这些方法表明解码策略强烈影响质量-效率权衡，但它们的接受边界通常是人工设计的，可能需要特定于机制的校准。最近，Learn2PD（Bao 等人，2026）提出了一种方法，保持基础 D-LLM 冻结，并使用与最终解码序列的一致性来训练一个轻量级令牌过滤器。TraceLock 共享从未来一致性学习的想法，但在所学习的策略上有所不同。TraceLock 不是学习一个与特定解码接口绑定的固定模块级过滤器，而是学习一个可变长度轨迹状态提交策略。我们的策略使用上下文隐藏状态、短程状态动力学和序列条件阈值，使其能够跨局部窗口、生成长度和步数预算进行泛化，无需改变架构或检查点。

## 3 方法

我们将 TraceLock 描述为一个插件式控制器，在冻结的 D-LLM 的解码循环内部实现一个学习到的令牌提交策略。基础 D-LLM 提出令牌和隐藏状态；TraceLock 决定哪些被提议的令牌应成为最终令牌。在控制层面，目标是学习一个策略，将当前轨迹状态映射为提交或修订决策。在训练过程中，完成的生成轨迹提供自监督的未来稳定性标签。在部署时，相同的学习控制器将此策略应用于当前轨迹状态，无需重新训练或针对每个设置进行校准。图 2(a) 展示了完成的轨迹如何定义用于训练的未来稳定性标签。图 2(b) 展示了部署时相应的在线决策：给定当前部分轨迹，控制器预测哪些活跃令牌足够稳定可以锁定。

（参见图标题）
图 2：TraceLock 概述。(a) 完成的扩散轨迹提供密集的令牌级监督：如果中间令牌与最终完成轨迹中相同位置的令牌一致，则被标记为正确。在部署时，TraceLock 预测相同的未来稳定性事件，以决定是否应接受被提议的令牌。这里的“正确”指的是与最终轨迹令牌的一致，而非任务级答案的正确性。(b) TraceLock 从当前轨迹状态对活跃令牌进行评分，并与序列条件动态阈值进行比较。软裁剪将提交决策限制在活跃的局部窗口内，而提示令牌和先前接受的令牌则作为上下文。

### 3.1 问题形式化

设一个提示占用位置 1,…,P，生成区域长度为 N，因此总序列长度为 L = P + N。一个掩码扩散语言模型迭代地更新一个序列：

x_t ∈ V^L, t=0,...,T,

其中 T 是总步数，x_t 是步骤 t 的预测令牌，V 是词汇表。未填充的生成位置包含一个特殊的掩码令牌。

在每一步，冻结的 D-LLM（参数为 ψ）产生令牌 logits 和隐藏状态：

(ℓ_t, H_t) = F_ψ(x_t),

其中 ℓ_{t,i} ∈ R^{|V|} 是位置 i 的令牌 logit 向量，H_t 是相应的内部隐藏表示。每个位置有一个状态，由下式表示：

s_{t,i} ∈ {prompt, gen, locked, eot}.

Prompt 位置是不可变的。Locked 位置是已经接受且不会被修订的已生成令牌。活跃的 gen 位置是可修订的生成位置，控制器对其施加作用。eot 状态标记生成的结束文本或填充位置，这些位置终止解码答案，在形成最终响应时由分词器移除。因此，核心决策是一个转换：

gen → locked.

给定当前候选令牌：

\hat{x}_{t,i} = arg max_{v∈V} ℓ_{t,i,v},

控制器决定是否提交该候选：

u_{t,i} ∈ {0,1}, i ∈ G_t,

其中 G_t = {i: s_{t,i}=gen}。如果 u_{t,i}=1，则该位置被填充为 \hat{x}_{t,i} 并变为 locked；否则它保持掩码状态以供后续精炼。

这种形式化将令牌提议与令牌提交分离。D-LLM 提出候选令牌；TraceLock 仅控制候选何时应停止被修订。可能的提交轨迹空间是组合的：在多个非空轮次中提交 N 个位置会产生指数级的可能性。穷举搜索该轨迹空间是不可行的，因此我们的目标是通过学习一个策略来摊销搜索，该策略倾向于那些已经与最终完成轨迹一致的提交。TraceLock 不会扩大基础模型的表达能力；它只是将解码偏向于基础模型可达轨迹空间内更有利的轨迹。

### 3.2 从未来稳定性学习提交策略

#### 自监督的未来稳定性标签。
对于每个完成的轨迹，我们观察到最终序列 x^⋆ = x_T 以及候选令牌