本文针对自回归(AR)大语言模型(LLM)固有的顺序解码导致推理效率低下的问题,提出了 Fast-dLLM v2,一个精心设计的块扩散语言模型(dLLM)。该模型旨在将预训练的 AR 模型高效地适配为 dLLM,以实现并行文本生成。
核心问题与研究目标
- 问题: 传统的 AR LLM 采用逐个 token 生成的方式,这种严格的从左到右的顺序限制了其在解码过程中的并行性,导致推理效率不高。
- 替代方案: 基于扩散的语言模型(dLLM)允许多个 tokens 联合预测或优化,理论上具有更高的解码并行度。然而,现有的 dLLM 存在无法有效使用 KV 缓存、推理延迟高、序列长度固定或灵活性受限等缺点。
- 研究目标: 提出一种块扩散语言模型(BD-LM),它结合了 AR 和扩散模型的优点,在块间自回归生成,在块内进行扩散。目标是设计一个能够扩展到大型 LLM 并在真实世界任务中保持高质量文本生成和强大扩展性的模型,同时实现比 AR 模型更快的推理速度。
创新点与主要贡献
1. 数据高效的后训练适配策略: 识别出本文块状注意力设计的 AR 友好特性,并利用此特性提出了一种后训练策略,仅需约 1B tokens 的微调数据即可将预训练 AR 模型适配为块扩散框架,而无需从头开始训练。相比于 Dream (Ye et al., 2025a) 需要约 500B tokens,数据效率提升了 500 倍,实现了无损适配。
2. 分层缓存与并行解码推理策略: 引入了一种结合了分层缓存机制和块状并行解码的推理策略。该设计包括一个用于跨块存储历史上下文表示的块级缓存,以及一个支持在部分解码的块内进行高效并行生成的子块缓存。这有效地重用了上下文信息并加速了块内 token 的生成。
3. 全面的大规模实验验证: 在高达 7B 参数的模型和多样化的任务上进行了广泛实验。结果表明,Fast-dLLM v2 在保持与 AR 基线相当的生成质量的同时,实现了比标准 AR 解码高达 2.5 倍的速度提升,并在 dLLM 中达到了最先进的效率。
基础定义。一个长度为 $L$ 的 token 序列表示为 $x = \{x_1, x_2, . . . , x_L\}$。传统的自回归模型通过建模条件分布 $p_\theta(x_i|x_{<i})$ 来顺序生成 $x$,并通过最小化交叉熵损失进行训练。</p>
扩散语言模型。扩散语言模型通过一个前向加噪过程和一个学习的反向去噪模型来定义生成分布。在时间 $t \in (0, 1)$,原始序列 $x_0$ 中的每个 token 以概率 $t$ 被独立地掩码(masked),产生一个损坏的序列 $x_t$。反向模型 $p_\theta(x_0|x_t)$ 则根据加噪后的输入预测原始的 tokens。
图2 | Fast-dLLM-v2 的训练过程。输入序列逐块解码。在每个块内,模型执行带有部分掩码的下一词元预测。为确保每个词元都得到训练,引入了互补掩码,使得一个视图中被掩码的词元可以从另一个视图中预测。我们只对以绿色高亮显示的预测词元应用损失,虚线连接了掩码词元及其对应的预测。
模型的训练目标是最小化期望的掩码 token 预测损失:
$$\mathcal{L}(\theta)=-\mathbb{E}_{t, x_{0}, x_{t}}\left[\frac{1}{t} \sum_{i=1}^{L} \mathbf{1}\left[x_{t}^{i}=[\mathrm{MASK}]\right] \log p_{\theta}\left(x_{0}^{i} \mid x_{t}\right)\right],$$其中 $t \sim \text{Uniform}(0, 1)$,$x_t$ 从前向过程中采样。损失仅在被掩码的 tokens 上计算。
基于预训练模型的微调框架。我们的块状扩散训练流程建立在预训练的 Qwen2.5-Instruct 模型(包括 1.5B 和 7B 版本)之上,具体实现为在指令调优数据上进行监督式微调(SFT)。每个训练批次都使用我们的块状扩散设置构建。具体而言,我们在每个块内引入部分 token 掩码,并结合互补掩码策略 (Li et al., 2025b),以确保每个 token 都能在可见和被掩码的上下文中得到训练。整体架构和训练工作流程如图 2 所示。
块状组织。给定一组 tokenized 样本,我们首先通过附加必要的 [MASK] token,将每个序列填充到其长度是块大小 $B$ 的整数倍。这些填充的 tokens 在损失计算中被忽略,且不参与梯度更新。完成对齐后,我们将填充后的序列拼接成一个长的 token 流,并将其分割成固定上下文长度 $C$ 的训练序列。因此,每个打包的序列被自然地划分为 $K = C/B$ 个大小为 $B$ 的不重叠块。这种块对齐的打包方式确保了高效的批处理,同时避免了块边界跨越样本边界。
带有互补视图的掩码令牌预测。对于每个块,我们随机采样一个二元掩码 $m \in {0, 1}^B$,其中 $m_i = 1$ 表示位置 $i$ 被一个学习到的 [MASK] 嵌入所替代。为了确保所有 tokens 在训练过程中都能同时接受到被掩码和未被掩码的监督,我们采用了互补掩码策略:每个训练样本被复制成两个视图,分别使用掩码 $m$ 和其补集 $\bar{m} = 1 - m$。这两个视图被放置在同一个批次中,这样模型可以同时观察到跨视图的掩码和未掩码上下文。
用于预测的令牌移位。为了保持预训练 AR 模型的表示质量 (Ye et al., 2025b,a),我们采用了一种移位标签(shifted-label)策略:在被掩码的位置 $i$ 预测 token 时,使用其前一个位置 $(i - 1)$ 的 logit。具体来说,如果 $x_t^i$ 被掩码,模型使用位置 $i - 1$ 的隐藏状态来预测 $x_0^i$,这与因果语言模型中的下一 token 预测机制保持一致。这使得 dLLM 在支持块内扩散的同时,能够维持类似 AR 的时序表示。
训练目标。我们最小化仅针对被掩码 token 的交叉熵损失:
$$\mathcal{L}_{\text{block}}(\theta) = -\mathbb{E}_{x,m} \left[ \sum_{i=1}^{L} \mathbf{1}[x_t^i = [\text{MASK}]] \log p_\theta(x_0^i \mid x_{<i}, x_{\text{block}(i)}) \right],$$ <p>其中,$x_{\text{block}(i)}$ 表示包含位置 $i$ 的块中的所有 tokens(包括掩码和未掩码的),$x_{<i}$ 是来自前面块的干净 tokens。块状注意力掩码。我们使用一种类似于 (Arriola et al., 2025) 中的混合注意力方案。对于每个训练样本,我们将加噪序列 $x_t$ 和其对应的干净序列 $x_0$ 沿序列维度拼接,总长度为 $2C$。然后应用注意力掩码 $M \in {0, 1}^{2C \times 2C}$ 来控制因果连接和双向连接。这种设计允许同时处理损坏和干净的上下文,从而促进互补掩码的监督。该注意力掩码自然地支持块并行性和块间的因果自回归依赖。我们进一步采用 flex-attention 实现来高效地实现这种结构化掩码并加速训练。
推理策略概述。在推理时,Fast-dLLM v2 采用一种块状解码策略,该策略平衡了 LLM 的自回归特性和扩散解码所提供的并行性。如图 3 所示,生成过程一次一个块地进行:先前解码的块被缓存并作为干净的前缀上下文重用,而当前块则进行并行的掩码 token 优化。
性能优势。通过结合块级缓存、块内并行解码和 DualCache 重用,Fast-dLLM v2 最大化了解码效率,而无需辅助模型或额外的推理开销。根据经验,我们观察到与标准 AR 解码相比,速度提升高达 2.5 倍,同时生成质量与强大的自回归基线相当。这使得 Fast-dLLM v2 成为在延迟敏感应用中实际部署 LLM 的一个有吸引力的选择。
带有缓存的块状自回归解码。由于 Fast-dLLM v2 中的每个块都按因果顺序解码,我们自然地保留了跨块的从左到右的语义。每个块解码后,其未掩码的 tokens 被缓存为只读上下文,供后续块使用。这种设计实现了块级的键值(KV)缓存重用,并显著减少了冗余计算。推理时的注意力掩码允许每个块在内部进行双向关注,同时对前面的块进行因果关注,这与训练期间的配置相呼应。
块内并行优化。为了加速块内的生成,我们采用了 Fast-dLLM (Wu et al., 2025) 中提出的基于置信度的并行解码策略。具体来说,我们根据模型对预测 tokens 的置信度,迭代地优化当前块中被掩码的 tokens。超过置信度阈值的 tokens 将被并行地解码和揭开,而不确定的位置则保持掩码状态以待后续优化。这避免了模糊的预测,并减少了生成延迟。
用于子块重用的 DualCache。为了在块内解码期间进一步减少冗余计算,我们集成了来自 Fast-dLLM 的 DualCache 机制。DualCache 为部分解码的块维护前缀和后缀的 KV 缓存,从而在揭示更多 tokens 时实现高效的重新计算。这种分层缓存不仅排除了昂贵的重新计算,还支持了基于置信度的优化中所使用的迭代式、选择性解码模式。
带填充的批量解码。为了支持对不同目标长度序列的批量生成,我们用 [MASK] tokens 对每个序列进行右填充,使其总长度能被块大小 $B$ 整除。然后将这些序列分组并逐块解码。在每一步中,批次中的所有序列都并行解码下一个块,无论还剩下多少真实的 tokens,从而确保在现代硬件上实现一致且高效的调度。
硬件配置:
软件配置:
训练超参数:
默认推理配置 (除非另有说明):
主要性能与速度结果
模型性能:
表1 | 各语言模型在一系列评估任务上的基准测试结果。模型按参数规模分为1B和7B+类别。评估指标包括代码生成(HumanEval、MBPP)、数学推理(GSM8K、MATH)、指令遵循(IFEval)、知识密集型问答(MMLU、GPQA)以及综合平均分(Avg.)。“Base”和“Plus”指使用EvalPlus进行代码基准测试的不同评估设置。每列的最佳结果以粗体显示,次佳结果带下划线。
速度与质量权衡:
吞吐量比较:
消融研究
训练策略影响 (表 2):
表2 | 不同 token 移位策略的基准测试结果。“+ CM”代表“+ 互补掩码”。每个基准测试的最佳性能以粗体显示,次佳性能带下划线。
子块大小与块大小的影响 (表 3 & 表 4):
表3 | 子块大小解码提升性能,大小为8时最优。
表4 | 使用不匹配大小的推理会降低性能。
子块缓存的影响 (图 6):
本文提出了 Fast-dLLM v2,一个可扩展的块扩散语言模型框架,它能够将预训练的自回归 LLM 高效地适配为用于并行文本生成的扩散式解码器。通过将块状扩散机制与互补掩码相结合,Fast-dLLM v2 在保留原始 AR 模型预测能力的同时,实现了块内的双向上下文建模。为了解决现有扩散模型的延迟问题,我们进一步提出了一种分层缓存策略,包括用于块间上下文重用的块级缓存和基于 DualCache 的用于块内高效优化的子块缓存,并结合了并行解码流程。在大型 Qwen2.5-Instruct 模型(1.5B 和 7B)上的广泛实验表明,Fast-dLLM v2 在不损失生成质量的情况下,实现了比标准 AR 解码高达 2.5 倍的速度提升,持续匹敌强大的 AR 基线,并在效率上超越了先前的扩散模型。这些结果凸显了块扩散框架作为在现实世界应用中部署高质量、低延迟 LLM 的一条可行路径的潜力。
基于 Qwen2.5-Instruct 的微调。我们在块状扩散训练框架下对预训练的 Qwen2.5-Instruct 模型(1.5B 和 7B)进行微调。除非另有说明,所有实验均采用 2048 的上下文长度和 256 的批量大小。训练在 64 个 NVIDIA A100 GPU 上使用 DeepSpeed Zero-3 进行。
训练数据。我们的模型在 LLaMA-Nemotron 后训练数据集的一个子集上进行微调,该数据集包含覆盖广泛领域的高质量指令遵循示例。我们使用块状打包预处理数据集,并对每个序列进行填充,使其长度成为块大小的倍数,以避免块边界错位。冗余的填充 token 不参与损失计算和梯度更新。
超参数。1.5B 模型使用 $2 \times 10^{-5}$ 的学习率训练 6,000 步,而 7B 模型使用 $1 \times 10^{-5}$ 的学习率训练 2,500 步。两种设置均使用 AdamW 作为优化器,并在前 500 步应用线性学习率预热。在上下文长度为 2048 和批量大小为 256 的情况下,每个训练步骤处理 524,288 个 tokens。这对应的总训练 token 数约为:
混合注意力掩码。为了在输入的损坏视图和干净视图之间实现高效和结构化的学习,我们使用了一种自定义的块感知注意力方案 (Arriola et al., 2025)。在每个训练步骤中,我们将加噪序列 $x_t$ 和干净序列 $x_0$ 拼接成一个总长度为 $2C$ 的单一输入,然后通过一个注意力掩码 $\mathcal{M}_{\text{full}} \in \{0, 1\}^{2C \times 2C}$ 应用混合注意力模式。
条件概率建模。为了简化表示,我们沿用先前工作的做法,略微滥用符号 $x_k$,在此上下文中它表示第 $k$ 个块中的 tokens 集合(而非早期章节中的第 $k$ 个 token)。具体来说,我们的目标是为所有块 $k \in [1, K]$ 建模条件概率 $p(x_k | x_t^k, x_{<k})$,其中 $x_t^k$ 是块 $k$ 的加噪版本,$x_{<k}$ 包含先前块中的所有干净 tokens。这种表述使我们能够通过将它们的拼接序列送入 transformer 并应用如图 7(a) 所示的精心构建的注意力掩码 $\mathcal{M}_{\text{full}}$,来同时处理加噪和干净的表示。</p>
掩码分解。整体注意力掩码可以分解为四个子掩码:
$$\begin{aligned} \mathcal{M}_{\text{full}} = \begin{bmatrix} \mathcal{M}_{BD} & \mathcal{M}_{OBC} \\ 0 & \mathcal{M}_{BC} \end{bmatrix}, \end{aligned}$$其中:
- $\mathcal{M}_{BD}$ (块对角掩码): 在加噪序列 $x_t$ 中,为同一块内的 tokens 提供双向自注意力,以实现块内优化:
- $\mathcal{M}_{OBC}$ (偏移块因果掩码): 允许 $x_t$ 中的每个加噪 token 关注干净序列 $x_0$ 中先前块的 tokens,从而保持块间的因果条件:
推理时的注意力机制。在推理时,我们采用一种简化的因果注意力机制,将已解码的块作为固定的前缀上下文重用。如图 7(b) 所示,先前生成的块 $x_{<k}^0$ 被缓存以避免冗余计算,只有当前的加噪块 $x_t^k$ 被主动优化。该块在内部进行双向关注,类似于训练时的 $\mathcal{M}_{BD}$,同时对先前块中的未掩码 tokens 进行因果关注。因此,注意力计算被限制在当前块及其因果前缀上,通过键值缓存重用和减少内存占用实现了高效解码。这种结构在保持跨块从左到右语义的同时,允许块内去噪并行进行。</p>
损失函数。我们最小化仅针对掩码 token 的交叉熵损失:
$$\mathcal{L}_{\text{block}}(\theta) = -\mathbb{E}_{x,m} \left[ \sum_{i=1}^{L} \mathbf{1}[x_t^i = [\text{MASK}]] \log p_\theta(x_0^i \mid x_{<i}, x_{\text{block}(i)}) \right].$$ <p>归一化系数的省略。值得注意的是,该目标函数似乎省略了标准掩码建模损失中常见的归一化系数 $1/L$(例如,除以掩码 token 的数量)。这是有意为之的,并由我们的互补掩码策略来证明其合理性。因为我们为每个训练样本 $x_0$ 使用一个互补掩码:我们总是采样两个互补的时间 $t$ 和 $1-t$,对应掩码 $m$ 和 $\bar{m} = 1 - m$。 $$-\left[\sum_{i=1}^{L} \mathbf{1}\left[x_{t}^{i}=[\mathrm{MASK}]\right] \log p_{\theta}\left(x_{0}^{i} \mid x_{<i}, x_{\mathrm{block}(i)}\right)\right]+\left[\sum_{i=1}^{L} \mathbf{1}\left[x_{1-t}^{i}=[\mathrm{MASK}]\right] \log p_{\theta}\left(x_{0}^{i} \mid x_{<i}, x_{\mathrm{block}(i)}\right)\right].$$ <p>由于互补掩码的存在,对于任何给定的样本 $x_0$,对损失有贡献的 token 总数始终是完整的序列长度 $L$。评估设置。我们在一系列涵盖推理、知识和代码生成的下游基准上评估所有训练好的模型。除非另有说明,所有评估均使用贪心解码(argmax)进行。我们对所有任务采用零样本(zero-shot)设置,但 GPQA 除外,该任务按照标准协议在 5 样本(5-shot)提示下进行评估。
评估框架与配置。所有非代码任务均使用 LM-Eval 工具包进行评估,以确保兼容性和公平的性能报告。对于 HumanEval 和 MBPP 等代码任务,我们采用 EvalPlus 框架进行可靠的通过率计算。除非另有说明,推理期间使用以下设置:
- 块大小 = 32
- 子块大小 = 8
- 并行解码禁用 (阈值 = 1)
该配置确保了训练和推理设置之间的一致性,有助于有效评估 Fast-dLLM v2 的块状扩散能力。
能力展示。为了更好地说明 Fast-dLLM v2 (7B) 的推理和交互能力,我们对单轮和多轮对话场景进行了详细的考察。代表性示例如表 5 和表 6 所示。
单轮对话场景。如表 5 所示,Fast-dLLM v2 能够处理单次交互中的复杂查询。在数学示例中,模型正确分析了有理函数以确定垂直渐近线的数量,应用了代数因式分解并求解了 $x$ 的未定义值。在代码部分,模型生成了一个正确且递归的 Python 实现来解决汉诺塔问题,并附有对输入参数和输出的适当解释。
多轮对话场景。表 6 突出了多轮对话,其中 Fast-dLLM v2 能够保留上下文并在前几轮的基础上进行构建。日常生活示例展示了模型进行时间推理的能力,例如计算公司成立以来的年数并确定其未来的周年纪念日。数学示例展示了逐步的逻辑推理,以解决一个涉及卡车载重量和总交付时间的现实世界算术问题。模型有效地将问题分解为顺序步骤,执行中间计算,并以要求的格式(小时和分钟)呈现最终结果。
综合能力。这些案例研究共同证明了 Fast-dLLM v2 在单轮和多轮设置中在数学推理、代码生成、时间理解和上下文连贯性方面的强大实力。
使用范围。在手稿准备过程中,我们使用了大型语言模型——但严格仅用于段落和句子的语言润色(语法、流畅性和语调)。这些工具未被用于产生想法、设计实验或确定结论。
作者责任。所有技术内容、方法论和解释均由作者撰写、验证和批准。为减少事实偏差或引用错误的风险,我们要求对模型编辑的每一句话进行人工审查,并对照原始来源核对所有参考文献。作者对手稿的准确性和完整性负全部责任。
数学
用户
函数 $y = \frac{2}{x^2+x-6}$ 的图像有多少条垂直渐近线?
助手
要确定函数 $y = \frac{2}{x^2+x-6}$ 的垂直渐近线数量,我们需要找出使分母为零的 $x$ 值,因为在这些点函数值趋于无穷大,这表明存在垂直渐近线。
首先,我们对分母 $x^2 + x - 6$ 进行因式分解。我们寻找两个数,它们的乘积为 -6,和为 1。这两个数是 3 和 -2。因此,我们可以将这个二次式分解为:
$x^2+x-6=(x+3)(x-2)$
所以函数可以重写为:
$$y = \frac{2}{(x+3)(x-2)}$$当 $x + 3 = 0$ 或 $x - 2 = 0$ 时,分母为零。解这些方程,我们得到:
$x=-3 \text{ and } x=2$
由于函数在 $x = -3$ 和 $x = 2$ 处无定义,所以在这些点存在垂直渐近线。
因此,函数 $y = \frac{2}{x^2+x-6}$ 有两条垂直渐近线。
垂直渐近线的数量是 2
代码
用户
编写一个递归的 Python 函数来解决汉诺塔问题并打印每一步的移动。
助手
def tower_of_hanoi(n, source, target, auxiliary):
if n > 0:
tower_of_hanoi(n - 1, source, auxiliary, target)
print(f"Move disk {n} from {source} to {target}")
tower_of_hanoi(n - 1, auxiliary, target, source)
tower_of_hanoi(3, 'A', 'C', 'B')
该函数使用递归解决了汉诺塔问题。它接受四个参数:n(盘子的数量)、source(起始柱)、target(目标柱)和 auxiliary(辅助柱)。函数会以“Move disk n from source to target”的形式打印每一步的移动。
表6 | Fast-dLLM v2 (7B) 的多轮对话案例。
数学
用户
一辆送货卡车一次可以运送 120 个箱子。如果一个仓库需要送出 960 个箱子,卡车需要跑多少趟?
助手
要确定送货卡车需要多少趟才能送出 960 个箱子,我们首先注意到卡车一次的运载能力。卡车一次可以运送 120 个箱子。我们需要计算运输 960 个箱子需要多少趟。
我们可以将这个问题设为一个除法问题:
$$\text{Number of trips} = \frac{\text{Total number of boxes}}{\text{Capacity of the truck per trip}}$$将给定值代入方程,我们得到:
$$\text{Number of trips} = \frac{960}{120}$$接下来,我们进行除法运算:
$$\frac{960}{120}=8$$因此,卡车需要 8 趟才能送出 960 个箱子。
用户
如果每趟需要 45 分钟,总共需要多长时间才能送完所有的箱子?请用小时和分钟作答。
助手
要确定送完所有箱子所需的总时间,我们需要遵循以下步骤:
步骤 1:计算所需的行程次数。
仓库需要送出 960 个箱子,而卡车每次行程可以运送 120 个箱子。因此,所需的行程次数是:
$$ \text{Number of trips} = \frac{960}{120} = 8 $$步骤 2:确定所有行程花费的总时间。
每次行程需要 45 分钟,所以 8 次行程的总时间是:
Total time $= 8 \times 45 = 360$ minutes
步骤 3:将总时间从分钟转换为小时和分钟。
一小时有 60 分钟,所以我们将 360 分钟转换为小时,通过除以 60:
$360 \div 60=6 \text{ hours}$
因此,送完所有箱子所需的总时间是:
6 小时