🌿 DFlash: Block Diffusion For Flash Speculative Decoding

大语言模型的推理主要采用逐 token 的顺序生成过程，其中，每个输出都依赖于此前的完整上下文。这种固有的串行性构成主要的性能瓶颈

• 推理速度缓慢、受限于内存带宽，并且无法充分利用现代 GPU。随着长思维链推理模型的出现，这一瓶颈变得愈发突出，因为耗时较长的推理如今已成为生成过程中的主要开销
• 投机解码已成为缓解这一瓶颈的主要方案。该范式使用轻量级草稿模型预测一段未来 token 序列，再由大型目标模型进行并行验证
• 尽管这一方法能够实现无损加速，并且已被广泛集成到生产环境中的推理框架，但 EAGLE-3 等先进方法仍然依赖自回归草稿生成。这种串行的草稿生成过程不仅本身效率低下，而且容易产生误差累积，从而使实际可达到的加速比受到约 2–3 倍的限制

近年来，扩散大语言模型通过支持并行文本生成和双向上下文建模，为自回归大语言模型提供一种有前景的替代方案

• 块扩散模型能够同时对一个块中的掩蔽 token 进行去噪。然而，当前开源扩散语言模型的生成质量通常不及相应的自回归模型。此外，为维持可接受的输出质量，往往需要较多的去噪步骤，这显著降低其实际推理速度
• 自回归模型具有更优的性能，却受到顺序生成延迟的限制；扩散模型支持快速的并行生成，但往往需要付出准确率下降的代价。由此自然引出一个研究问题：能否结合这两种范式的优势，同时缓解它们各自的不足？ 一种颇具吸引力的解决方案是，利用扩散模型进行高速、并行的草稿生成，同时依靠高质量的自回归模型进行验证，以确保最终输出保持无损

然而，将扩散模型用于草稿生成并非易事，现有方法要么缺乏实用性，要么只能获得有限的加速收益

• DiffuSpec 和 SpecDiff-2 等方法使用规模庞大的草稿模型，例如参数量达到 7B 的模型。如此高的内存占用，对实际模型服务而言往往难以承受
• 尽管这些大型草稿模型能够生成质量较高的草稿 token ，并获得较长的接受长度，但较高的草稿生成延迟将其实际加速比限制在较为有限的 3–4 倍。相比之下，PARD 训练小型自回归模型，使其模仿扩散式并行生成，再利用这些模型对目标大语言模型执行投机解码。然而，得到的小模型不具备目标大语言模型的建模能力，因此接受长度有限，加速比的上限约为 3 倍

本文提出 DFlash，一种利用轻量级块扩散模型，同时实现快速、高质量草稿生成的投机解码框架

• 关键认识十分简单：目标模型最清楚。 正如 Samragh et al. 所观察到的，大型自回归语言模型的隐藏特征隐式包含多个未来 token 的信息。DFlash 将这些隐藏特征用作上下文，以此作为条件，使草稿模型能够并行预测未来的 token 块
• 实际上，草稿模型成为一个扩散适配器，能够高效利用大型目标模型所建模的深层上下文特征。DFlash 不再要求一个微型草稿模型从头进行推理，而是将目标模型的推理能力与小型扩散草稿模型的并行生成速度结合起来

🍄 Related Work

投机解码通过缓解自回归生成中的顺序执行瓶颈，加速大语言模型推理。早期方法使用较小的草稿模型生成 token 序列，再由较大的目标模型进行并行验证

• Medusa 通过为基础大语言模型增加多个预测头，并使用树形注意力进行并行验证，消除对外部草稿模型的依赖。EAGLE 系列进一步利用冻结目标模型提供的特征级上下文，改进投机解码。EAGLE-1 通过预测未来隐藏状态的分布来提高接受率，EAGLE-2 引入自适应草稿树，EAGLE-3 则改进训练目标，使加速收益能够随规模扩展而提升
• 尽管取得上述进展，大多数现有方法仍然依赖自回归草稿生成。这一过程在本质上仍然是顺序执行的，因此限制其加速效果

扩散大语言模型通过并行预测被掩蔽的 token ，为自回归生成提供一种替代方案。LLaDA 首次将扩散大语言模型扩展至数十亿参数规模，取得与 LLaMA-3.1-8B相当的性能。然而，完全并行的扩散模型受限于固定长度生成，并且缺乏高效的 KV 缓存支持

• 块扩散模型通过逐块对序列进行去噪，将并行性与自回归结构相结合，从而解决上述问题。在此基础上，Fast-dLLM v2 和 SDAR 将预训练的自回归大语言模型适配为块扩散变体，在特定任务上保持生成质量的同时，实现并行生成
• 尽管如此，现有扩散大语言模型的整体性能通常仍不及先进的自回归模型，而且往往需要较多的去噪步骤，从而限制其实际推理速度

近期研究开始探索在投机解码中使用扩散模型作为草稿模型。TiDAR 联合优化扩散目标与自回归目标，通过扩散实现并行思考，通过自回归解码实现顺序表达，但尚不能无损地保持最终生成质量

• 另一些方法则将自回归模型改造为扩散式草稿生成模型。Samragh et al. 观察到，自回归大语言模型隐式编码未来 token 的信息，并通过训练 LoRA 适配器实现并行草稿生成，同时保留基础模型用于验证。
• DiffuSpec 和 SpecDiff-2 使用大型预训练扩散语言模型作为投机草稿模型，并通过推理时搜索或训练—测试对齐来提高接受率。然而，这些方法依赖规模庞大的草稿模型，例如参数量达到 7B 的模型，因此会引入显著的内存和延迟开销。尽管它们能够获得较长的接受长度，但高昂的草稿生成成本往往会抵消其在真实模型服务场景中的实际加速收益

🌊 Inference

来自目标模型的上下文特征：An et al. 等先前研究直接将小型扩散模型用作投机草稿模型，导致接受长度较短，加速收益有限。为验证这一点，训练一个不使用任何目标模型条件信息的五层块扩散草稿模型，并在多个数学基准上进行评估。其加速收益较为有限，通常约为 2–3 倍。这一局限源于缺乏丰富的上下文引导：在无法访问目标模型内部表示的情况下，扩散草稿模型实际上必须从头预测未来 token

• 相比之下，大型自回归目标模型的隐藏表示所编码的信息，远比 token 层面的 logits 丰富。这些特征能够捕获长距离依赖关系和任务相关的语义信息；更重要的是，它们隐式编码有关未来 token 预测的信息
• 在 DFlash 中，给定输入提示，目标模型首先执行一次标准的预填充前向传播，生成第一个 token 。在这次前向传播中，从一组由浅至深均匀选取的固定层中提取隐藏表示。随后，将这些隐藏状态拼接，并通过一个轻量级投影层，将跨层信息融合为紧凑的目标上下文特征，用于为草稿模型提供条件信息

通过 KV 注入，使接受长度能够随模型深度扩展而提升：EAGLE-3 等现有方法同样利用目标模型的隐藏特征，但它们将这些特征与草稿模型的 token 嵌入融合，仅作为草稿模型的输入。随着草稿模型深度增加，来自目标模型的信息逐渐被稀释，导致增加草稿层数所带来的接受长度收益不断减小

• DFlash 将融合后的目标上下文特征视为持续保留的上下文信息，将其直接注入草稿模型每一层的键和值投影中。投影后的特征存储在草稿模型的 KV 缓存中，并在各轮草稿生成之间复用。这一设计在草稿模型的整个计算过程中提供充分且一致的条件信息，使接受长度能够随着草稿层数的增加而有效提升

并行扩散草稿生成：较低的草稿生成延迟是 DFlash 实现加速的另一项关键因素。自回归草稿模型必须执行多次顺序前向传播，才能生成草稿 token 或草稿树，这限制并行性，并导致 GPU 利用效率低下。相比之下，DFlash 通过块级扩散过程预测下一个 token 块。在一次前向传播中，块内所有被掩蔽的位置均被并行解码。与自回归草稿生成相比，这种逐块并行生成显著降低草稿生成延迟，并且即使使用更深的草稿模型，也能实现显著更高的硬件利用率

总体而言，DFlash 将基于扩散的并行草稿生成，与来自目标模型的紧密耦合条件信息相结合，在显著降低草稿生成延迟的同时，实现高质量的草稿生成

DFlash 草稿模型的训练目标，是使块级扩散预测与冻结的自回归目标模型的输出对齐。并未直接采用标准的块扩散训练方法，而是引入若干关键修改，以提高训练效率、可扩展性，以及与推理阶段投机解码行为的一致性

KV 注入：遵循推理流程，给定由提示及其响应组成的序列，首先将完整的干净序列输入目标模型，提取并融合所有 token 的隐藏特征。随后，以键和值投影的形式，将这些隐藏特征注入草稿模型

掩蔽块的随机采样：在标准块扩散训练中，响应被均匀划分为多个块，各块内部的随机位置被掩蔽，模型通过对这些掩蔽 token 进行去噪来训练

• 相比之下，DFlash 针对投机解码场景专门设计块的构造方式。从响应中随机采样锚点 token ，将每个锚点作为一个块的首个位置，并将其余位置掩蔽。草稿模型被训练为并行预测后续的个 token 。这与推理时的行为直接对应：草稿模型始终以目标模型生成的一个干净 token 为条件，即上一轮验证产生的额外 token （bonus token）。随机化锚点位置，还使草稿模型能够接触更加多样的目标上下文特征，从而提高数据利用效率和覆盖范围
• 训练期间，所有块被拼接为一个序列，并使用稀疏注意力掩码进行联合处理。 token 可以在同一块内部进行双向注意力交互，并访问相应的、通过注入提供的目标上下文特征，但不同块之间的注意力交互被禁止。借助 Flex Attention，这一设计能够在一次前向传播和反向传播中，高效训练多个草稿块

高效的长上下文训练：对于 EAGLE-3 等方法，由于其训练时测试（training-time test）开销较高，在长上下文上训练投机草稿模型具有较大挑战。DFlash 通过固定每条序列的掩蔽块数量，并在每个训练轮次中为各条序列随机采样锚点位置，实现高效的长上下文训练。这一策略在提供有效数据增强的同时，将训练开销控制在有界范围内

通过损失加权加快收敛：在投机解码中，并非所有 token 都同等重要。草稿块中较早位置的预测错误，会使其后的全部 token 失效。因此，较早位置的预测对接受长度具有更大的影响。通过对交叉熵损失进行加权，在训练时强调靠前的 token 位置，以体现这种不对称性

• 具体而言，对于块内位置为 的 token ，采用指数衰减权重

• 其中，控制衰减速率。这种加权方式优先强调较早的位置，能够加快收敛，并获得比均匀加权更长的接受长度

共享嵌入层与语言模型输出头：为提高训练效率，草稿模型与目标模型共享 token 嵌入层和语言模型输出头，并在训练期间将二者保持冻结。训练时，仅更新草稿模型的 Transformer 层。这一设计减少可训练参数的数量，并促使草稿模型充当一个轻量级扩散适配器，与目标模型的表示空间紧密对齐