生存引导的长度控制以实现高效扩散语言模型

# 面向高效扩散语言模型的生存引导式长度控制

来源：https://arxiv.org/html/2608.26374
Abbas Ghaddar¹ Yufei Cui  
机构：华为诺亚方舟实验室，蒙特利尔研究中心，加拿大  
邮箱：[{\{ivan.kobyzev,abbas.ghaddar,yufei.cui\}@huawei.com}](mailto:)

###### 摘要

扩散语言模型（\(DLMs\)）通过迭代去噪掩码序列来生成文本，但标准解码要么固定序列长度，要么依赖临时停止规则，常导致不必要的去噪步骤。我们将长度选择重新定义为序列结束符的离散时间生存问题，并提出一种可插入现有任意\(DLM\)的、免训练的长度预测器。在推理和代码生成基准测试中，生存引导式长度解码在保持任务准确性的前提下，将推理速度提升最高达\(7\times\)。我们进一步发现，即使在同一数据集中，预测长度也存在显著差异，这使得模型性能对所选长度高度敏感。

## 1 引言

掩码扩散语言模型（\(DLMs\)）[Austin et al. (2021a)](https://arxiv.org/html/2608.26374#bib.bib6); [Shi et al. (2024)](https://arxiv.org/html/2608.26374#bib.bib7)通过迭代去噪掩码画布来生成文本。在每个扩散步骤中，掩码位置根据模型的预测分布更新，逐步将全-\[MASK\]后缀转化为文本。这种迭代精炼实现了灵活的任意顺序解码和并行化，但也给解码成本带来一个核心问题：何时停止去噪？在每个步骤中提交哪些token？

掩码\(DLM\)的标准解码通常遵循任意顺序自回归（\(AOAR\)）模式 [Ou et al. (2024)](https://arxiv.org/html/2608.26374#bib.bib8)：选择一个与任务无关的大最大长度\(L_{\max}\)，运行固定的去噪调度，仅在所有掩码被移除或达到迭代预算时停止。虽然简单，但这种方法常常造成浪费：许多提示所需token远少于\(L_{\max}\)，但解码器仍在为不必要的长画布对应位置花费计算资源进行精炼。

因此，掩码\(DLM\)推理的主要开销并非去噪规则本身，而是保守的全局长度预算与实例特定所需长度之间的错配。本文聚焦于长度选择：对于给定提示，应生成多少个新token？我们证明长度选择可自然地用生存分析解释 [Andersen et al. (1993)](https://arxiv.org/html/2608.26374#bib.bib9)。

具体而言，我们将生成长度视为序列结束符 \[EOS\] 的离散时间生存变量。利用\(DLM\)在长掩码画布上单次前向传播得到的每位置 \[EOS\] 概率，我们获得了离散时间风险的插件估计，并通过标准生存恒等式恢复了预期长度的闭式估计。这产生了一个**免训练、与模型无关的长度预测器**，可插入现有掩码\(DLM\)中，无需修改模型参数或改变底层去噪调度。实证表明，生存引导式长度预测减少了序列真实结束位置之外的不必要精炼，在保持与使用足够大\(L_{\max}\)解码的基线性能相当的前提下加速了推理。

对两个大规模\(DLM\) [LLaDA](https://arxiv.org/html/2608.26374#bib.bib4) 和 [Dream](https://arxiv.org/html/2608.26374#bib.bib5) 的广泛实验一致表明，在推理和代码生成基准测试中，推理速度提升最高达\(7\times\)，且任务性能无损失。此外，我们的分析表明长度预测确实依赖于样本，即使在同一分布内，长度也随样本显著变化。

## 2 方法

生存分析 [Andersen et al. (1993)](https://arxiv.org/html/2608.26374#bib.bib9) 研究随机事件时间的分布，经典场景如故障时间或死亡时间。核心对象是风险函数，它建模在事件尚未发生的前提下，事件在时间\(t\)发生的概率。这一观点本质上是离散和序列的：在每个步骤中，我们要么存活至下一步，要么终止。将此框架应用于\(DLM\)，未知的序列长度成为序列结束符的离散事件时间，在每个候选位置我们询问：如果序列尚未结束，在此处结束的可能性有多大？

这使我们能够重用标准生存恒等式，从\(DLM\)的一次传播中获得预期长度的插件估计。考虑通过迭代去噪掩码画布生成文本的掩码扩散语言模型 [Austin et al. (2021a)](https://arxiv.org/html/2608.26374#bib.bib6); [Shi et al. (2024)](https://arxiv.org/html/2608.26374#bib.bib7)。设

\[ \mathbf{x}^{(0)} = [x_{1:P}, \texttt{[MASK]}, \dots, \texttt{[MASK]}] \tag{1} \]
为初始序列，包含长度为\(P\)的提示和\(T\)个掩码位置，记 \(\mathbf{x}^{(s)}\) 为\(s\)步去噪后的序列。在每个步骤\(s\)，\(DLM\)产生logits和概率：
\[ \mathbf{z}^{(s)} = f_\theta(\mathbf{x}^{(s)}, t_s), \quad p_{i,v}^{(s)} = \text{softmax}_v z_{i,v}^{(s)} \]
其中\(t_s\)是（离散）扩散时间。

我们将未知序列长度视为生成跨度位置上的离散时间生存问题。考虑如公式(1)所示的掩码输入，提示占据位置\(1, \dots, P\)，模型任务是在位置\(P+1, \dots, P+T\)生成最多\(T\)个新token。然而，理想生成长度未知。标准做法中\(T\)被设为与任务无关的大上界\(L_{\max}\)，导致许多提示在\(L_{\max}\)前早已终止，但解码器仍在所有\(T\)个槽位上运行完整的去噪调度。这种固定全局预算与理想实例特定长度之间的错配是计算浪费的主要来源。

为预测理想生成长度\(L\)，我们对初始画布运行一次扩散模型，并将所得logits记为\(\mathbf{z}^{(0)}\)。一步\(DLM\)解码后文本质量通常极低，但这些logits仍有信息量，可通过将其每位置 \[EOS\] 概率解释为长度生存信号用于我们的长度预测任务。

设\(t\)为logits \(\mathbf{z}^{(0)}\) 在序列中的位置，定义
\[ p_t := \mathbb{P}\big(\texttt{[EOS] at position } t \,\big|\, \text{prompt, masks}\big) = \mathrm{softmax}(\mathbf{z}^{(0)}_t)[\texttt{[EOS]}], \tag{2} \]
其中\(t = P+1, \dots, P+T\)。

我们用离散时间生存过程建模序列终止。定义相对位置\(k = 1, \dots, T\)的风险为在序列尚未终止的条件下，序列在此时终止的概率：
\[ h_k := \mathbb{P}(L=k \mid L \geq k, \text{prompt}). \]
假设扩散模型训练良好并拟合数据分布，可估计风险为 \(h_k \approx p_{P+k}\)。

我们采用平均场近似 [Blei et al. (2017)](https://arxiv.org/html/2608.26374#bib.bib10)，将给定提示下的各位置 \[EOS\] 事件视为条件独立。在此假设下，存活至步骤\(k\)的生存函数为
\[ S(k) = \prod_{i=1}^{k} \bigl(1 - h_i\bigr), \quad S(0)=1, \tag{3} \]
恰好在\(k\)终止的概率为
\[ \pi_k = \mathbb{P}(L=k) \approx h_k \, S(k-1). \tag{4} \]

给定插件长度分布 \(\{\pi_k\}_{k=1}^T\)，可计算截断期望长度：
\[ \mathbb{E}[L] = \sum_{k=1}^{T} k \, \pi_k = \sum_{k=1}^{T} S(k-1), \tag{5} \]
其中第二个等式是期望等于生存函数之和的标准恒等式（证明见附录A）。

测试时，我们选择单个插件长度估计 \(\hat{L}(x) = \mathbb{E}[L]\)，并将其作为最大新token数传递给\(DLM\)。关键的是，这不需要额外训练或参数。长度预测器完全从基础模型的 \[EOS\] logits推导而来。详见附录B的算法1。

## 3 实验

### 3.1 实验设置

我们实验使用两个以显著不同方式训练的强大掩码扩散语言模型（\(DLMs\)）： [LLaDA-8B-Base](https://arxiv.org/html/2608.26374#bib.bib4) 和 [Dream-v0-Base-7B](https://arxiv.org/html/2608.26374#bib.bib5)。两者均基于PyTorch [Paszke et al. (2019)](https://arxiv.org/html/2608.26374#bib.bib2) 和 Transformers库 [Wolf et al. (2020)](https://arxiv.org/html/2608.26374#bib.bib3) 实现。下文分别称这两个模型为LLaDA和Dream。

我们考虑在原模型评估中使用的一组推理和代码生成基准，包括 [BBH](https://arxiv.org/html/2608.26374#bib.bib18)、[GSM8K](https://arxiv.org/html/2608.26374#bib.bib19)、[MATH](https://arxiv.org/html/2608.26374#bib.bib20)、[HumanEval](https://arxiv.org/html/2608.26374#bib.bib21) 和 [MBPP](https://arxiv.org/html/2608.26374#bib.bib22)。遵循标准LM评估工具 [Gao et al. (2024)](https://arxiv.org/html/2608.26374#bib.bib1) 设置，报告BBH的3-shot准确率、GSM8K的严格匹配5-shot准确率、MATH的4-shot准确率、HumanEval（0-shot）的pass@1，以及MBPP（3-shot）的pass@1。使用每个模型的自定义评估代码和每基准超参数。LLaDA和Dream的作者均使用 \(L_{\max}=1024\) 的后缀掩码序列长度处理所有基准。所有实验在单个现代计算加速器上进行，报告批次大小为1的结果。

### 3.2 主要结果

表1：基线（w/o）与长度预测（w/）的解码速度（秒/样本）。括号内为相对于基线的加速比。
表1显示LLaDA和Dream在使用固定最大长度的标准AOAR解码（w/o）与装备我们的生存引导长度预测器（w/）时的解码速度（每样本秒数）。观察到，对两个模型，使用预测长度在所有基准上均带来显著解码速度提升，加速比从\(3.2\times\)到\(6.6\times\)不等。此外，我们发现除HumanEval外，尽管结构差异，大多数基准上两个模型的加速比范围相似。

表2：基线（w/o）与长度预测（w/）的任务性能及标准差。表2显示使用与不使用长度预测器的模型各任务性能及标准差。可见表1报告的效率提升并未以任务性能为代价。所有基准上，w/o与w/之间的差异均在报告的标准差范围内，表明无统计显著变化。

这表明我们的生存引导长度预测器可靠地修剪了冗余尾部步骤，同时基本保持生成序列质量不变。值得注意的是，该过程完全与模型无关且免训练：包括在长掩码画布上单次前向传播、将每位置 \[EOS\] 概率解释为离散时间风险、以及从生存曲线计算闭式预期长度。尽管架构和训练差异，此插件估计器一致带来\(3\times\)至\(7\times\)加速且无精度损失，支持生存引导长度控制作为扩散语言模型高效解码稳健基元的观点。

### 3.3 固定平均长度消融

表3：报告每任务平均预测长度\(\bar{L}\)及其标准差，以及使用\(\bar{L}\)作为固定范围时的性能变化\(\Delta\)及其标准差，与我们按样本生存引导长度选择相比。
为检验我们的增益是否仅源于选择较短的全局预算，而非实例级适应，我们运行一项消融实验：每个模型使用*单一*固定生成长度，等于我们方法得到的全数据集平均预测长度\(\bar{L}\)。表3报告每个模型和基准的平均预测长度\(\bar{L}\)及其标准差，并显示将按样本长度预测替换为固定范围\(\bar{L}\)（而非使用最大长度\(L_{\max}\)）时的性能变化\(\Delta\)。观察到基于平均预测长度的固定长度策略在大多数设置下性能不及生存引导解码，少数例外（如LLaDA在MBPP上）。我们归因于每个数据集内所需解码长度存在显著的样本间变异，即使在精心策划的标准基准中亦然。此观察表明按样本长度预测对维持性能很重要，因为两个模型似乎对所选解码范围敏感。更多结果和消融见附录C。

## 4 相关工作

#### 扩散语言模型。

文本离散去噪扩散模型首次出现在离散状态空间结构扩散工作中 [Austin et al. (2021a)](https://arxiv.org/html/2608.26374#bib.bib6)，后专门化为扩散式语言模型（\(DLMs\)），如LLaDA和Dream，它们在固定画布上迭代揭示token [Nie et al. (2025)](https://arxiv.org/html/2608.26374#bib.bib4); [Ye et al. (2025)](https://arxiv.org/html/2608.26374#bib.bib5)。这些模型以并行精炼和灵活token顺序换取了自回归LM严格的从左到右分解，但典型解码器仍假设保守的固定最大长度\(T\)和固定精炼步数。近期工作主要关注提升生成质量和扩展至更大模型，对长度预测和免训练解码策略的效率关注较少。

#### 扩散LM的长度控制。

与我们的长度预测器最接近的是 [DAEDAL](https://arxiv.org/html/2608.26374#bib.bib11)，它也针对\(DLM\)的固定长度限制。DAEDAL从