训练扩散模型进行从左到右推测

# 引导扩散模型从左到右进行推测

来源: https://arxiv.org/html/2606.11552

###### 摘要

大型语言模型（LLM）在广泛的任务中表现出色，但其自回归解码过程由于固有的顺序词元生成，导致推理成本高昂。推测解码通过使用轻量级草稿模型提出多个未来词元，然后由更大的目标模型并行验证，从而解决了这一瓶颈。近期工作表明，扩散语言模型非常适合此场景，因为它们可以并行生成整块草稿词元，从而缓解自回归草稿的顺序约束。这种机制的一个微妙之处在于，块扩散草稿生成器在块内**双向**生成词元，而验证则由自回归目标模型执行，该模型严格以**从左到右**的方式评估词元，导致对称的训练时目标与不对称的验证时奖励之间存在差距。在这项工作中，我们实证分析了三种弥合这一差距的训练时干预措施：词元位置加权、针对每个块内打破已接受前缀位置的首次错误焦点损失，以及用可微替代项替代期望接受长度的链损失项。这三种干预措施沿正交轴（位置、块条件首次错误、联合前缀）发挥作用，并具有加性复合效应；它们同样与测试时对齐机制（例如多草稿自我选择）正交，原则上可以与之结合。在四个目标模型和六个推理、代码及对话基准上，这三种干预措施相较于位置统一的基线，将每个基准的接受草稿长度提高了21%–76%，且无需增加额外的前向传播，也无需改变推理流程或拒绝采样的精确性契约。

参照图注

图1: 基于扩散的推测草稿模型中的训练-验证不匹配以及我们分析的三种训练时干预措施。(a) 标准训练。草稿模型在\(K\)词元块上使用完全双向注意力进行训练，因此每个位置对称地以所有其他位置为条件。然而，在推理时，目标模型严格从左到右验证该块：单个早期拒绝（此处为位置\(2\)）会截断整个后缀。草稿模型在位置\(3\)和\(4\)上的预测与真实情况abcd匹配，但这两个词元连同被拒绝的位置\(2\)词元一起被丢弃。(b) 对齐训练。我们分析了三种互补的训练时干预措施，使草稿模型的目标与因果接受契约更一致：位置级损失衰减（Weighting）、针对每个块中首个错误预测位置的辅助交叉熵项（Focal），以及联合前缀接受概率的可微替代项（Chaining）。它们共同提高了期望接受前缀长度，并恢复了双向训练已经提供的逐词元正确性。

## 1 引言

大型语言模型（LLM）已成为从对话助手和代码生成 (Chen et al. 2021; Rozière et al. 2023) 到检索增强问答、多步智能体系统以及长程推理 (OpenAI 2024; DeepSeek-AI et al. 2025) 等应用的核心计算基元。随着这些系统日益成熟，其生命周期中的主要成本已从训练转向推理：模型只训练一次，但持续提供服务，而每次查询的成本会因输出长度、递归智能体调用以及近期推理模型发出的长思维链而放大。因此，降低LLM推理的延迟和每词元成本是当代机器学习中最重要的系统问题之一。

单流LLM解码的瓶颈在于内存带宽而非算术吞吐量：每个自回归步骤将整个参数集从高带宽内存（HBM）流式传输以生成单个词元，这种不平衡随着每一代硬件的计算速度超过HBM速度而加剧 (Pope et al. 2023)。标准的应对措施——量化 (Dettmers et al. 2022; Frantar et al. 2023; Lin et al. 2024)、稀疏性 (Frantar and Alistarh 2023)、蒸馏 (Hinton et al. 2015) 以及系统级技术如FlashAttention (Dao et al. 2022)、PagedAttention和连续批处理 (Kwon et al. 2023; Yu et al. 2022)——都降低了单次前向传播的成本。而推测解码 (Leviathan et al. 2023; Chen et al. 2023) 则减少了所需**串行**前向传播的次数：一个廉价的草稿模型提出\(K\)个候选词元，目标模型在单个并行前向传播中验证它们，并通过拒绝采样步骤接受与目标分布一致的最长前缀。该过程是精确的——样本在统计上与目标模型完全一致——从而消除了其他技术中常见的质量与速度的权衡，使得核心问题纯粹在于草稿模型设计：如何构建一个既评估成本低廉又能良好对齐目标分布的草稿分布。

EAGLE系列 (Li et al. 2024b, a, 2025b) 通过在**特征空间**而非词元空间进行草稿生成来追求这一目标，将一个轻量级辅助头耦合到目标模型自身的隐藏状态。后续迭代引入了动态树扩展和多层特征聚合，EAGLE-3已成为生产环境中推测解码的事实标准基线。然而，EAGLE草稿模型在每个推测步骤内部仍然是**自回归**的，需要顺序生成词元预测，从而限制了其可实现的最大加速比。在这项工作中，我们专注于**块扩散草稿模型**——其中DFlash (Chen et al. 2026) 是近期一个突出的例子——这类模型用**扩散式**并行块解码器取代了自回归特征预测。在可配置的目标模型中间隐藏状态集条件下，这种草稿模型在一次非自回归前向传播中发出整个\(K\)词元块；目标模型在标准的拒绝采样契约下验证该块，保证了精确性。当\(K=16\)时，这使每步的上限比EAGLE-3提高三倍以上，同时通过多层条件作用保持了有竞争力的每词元接受率。

将块扩散草稿模型与自回归目标模型配对也暴露出一种在自回归草稿模型机制中不那么明显的表征不对称性：扩散模型在完全双向注意力下训练去噪块，因此每个位置都从两个方向获取上下文，而目标模型严格从左到右验证，只接受最长的因果一致前缀。因此，草稿模型必须不对称地分配其预测能力——早期位置承担了不成比例的重担，因为单次早期分歧会截断所有后续草稿词元，无论其质量如何——这一目标与扩散公式继承的对称去噪损失相冲突。本研究通过分析三种沿正交轴重塑位置级损失分布的训练时干预措施，实证地研究了这种张力；该分析与为同一类草稿模型开发的测试时对齐机制（例如ddTree (Ringel and Romano 2026)）是互补的。

## 2 相关工作

#### 高效LLM推理。
一系列广泛的技术减轻了单流解码的内存带宽成本：低位量化 (Dettmers et al. 2022; Frantar et al. 2023; Lin et al. 2024; Xiao et al. 2023)、稀疏性 (Frantar and Alistarh 2023; Sun et al. 2024)、蒸馏 (Hinton et al. 2015; Sanh et al. 2019)、FlashAttention (Dao et al. 2022; Dao 2024) 以及系统级技术如PagedAttention、连续批处理和分块预填充 (Kwon et al. 2023; Yu et al. 2022; Agrawal et al. 2023)。

#### 推测解码。
块并行解码最初被探索为一种确定性加速技术 (Stern et al. 2018)；现代保分布公式由 Leviathan et al. (2023) 和 Chen et al. (2023) 同时引入。SpecInfer (Miao et al. 2024) 将线性链草稿推广到在树结构注意力下并行验证的树，显著提高了期望接受长度。后续工作沿着大致正交的轴扩展了设计空间：Medusa (Cai et al. 2024) 将并行预测头直接附加到目标模型上，Lookahead Decoding (Fu et al. 2024) 通过基于n-gram轨迹的Jacobi式不动点迭代绕过了草稿模型训练，自推测解码 (Zhang et al. 2024a) 重用目标模型的一部分层作为草稿模型，在线推测解码 (Liu et al. 2024b) 使草稿模型持续适应部署的工作负载——每种方法都在不同的廉价候选机制下接受相同的精确性契约。

#### 特征级草稿模型。
一条特别成功的路线是从目标模型自身的中间表示而不是独立训练的小型LM中抽取候选。EAGLE (Li et al. 2024b) 引入了特征级自回归；EAGLE-2 (Li et al. 2024a) 添加了上下文相关的动态树扩展；EAGLE-3 (Li et al. 2025b) 用多层特征聚合增强了草稿模型，建立了生产环境中推测解码的事实标准基线。相关设计包括Hydra (Ankner et al. 2024)、GliDe with CaPE (Du et al. 2024)、Kangaroo (Liu et al. 2024a) 和 HASS (Zhang et al. 2024b)。所有这些在每个推测步骤内仍然是**自回归**的，因此它们的经验范围受限于树深度，而非单次草稿模型前向传播的成本。

#### 扩散和非自回归语言建模。
并行块生成在推测解码之外有着悠久的历史：非自回归翻译 (Gu et al. 2018) 首次展示了单次序列生成，迭代精炼 (Lee et al. 2018) 与诸如Mask-Predict (Ghazvininejad et al. 2019) 和 SUNDAE (Savinov et al. 2022) 等掩码LM解码器通过重复去噪弥补了大部分由此产生的质量差距。离散扩散LM (Austin et al. 2021a; Li et al. 2022; Gong et al. 2023; Gulrajani and Hashimoto 2023; Lou et al. 2024; Nie et al. 2025) 通过掩码块的双向去噪将这一范式扩展到LLM规模，而块扩散 (Arriola et al. 2025) 通过结合块内双向去噪与块间因果条件作用，在该范式和自回归之间进行插值。更近期的工作将扩散LM视为其自身的高效生成器：EfficientDLM (Fu et al. 2026b) 使用位置相关的词元掩码；TiDAR (Liu et al. 2025) 将单个模型同时充当草稿模型和验证器；Nemotron-Labs-Diffusion (Fu et al. 2026a) 引入了一种统一自回归、扩散和自推测解码的三模态架构。本文追求的自然综合——将块扩散模型作为推测草稿模型，其中DFlash (Chen et al. 2026) 是一个突出的近期例子——继承了这些模型的并行块加速能力，同时正如我们在第4节中所示，引入了一个与因果接受契约相冲突的训练目标。

#### 使草稿模型与接受契约对齐。
快速增长的文献针对我们所关注的同一差距。Zhou et al. (2024) 研究了自回归草稿模型的标准前向KL蒸馏损失的reverse-KL和总变分替代方案，而Samarin et al. (2026) 最近的LK损失优化了基于TV的每词元接受目标。这些方法改变了每词元的**损失族**，但均匀加权所有\(K\)个草稿位置，并且不关心接受契约的联合前缀结构；它们也是针对自回归草稿模型（Medusa式头、EAGLE变体、MTP模块）而非双向块解码器开发的。Sandler et al. (2025) 的SpecDiff-2 通过**streak蒸馏**——一种微调目标，在验证器采样的教师轨迹下最大化期望接受streak的可微替代项——并配合测试时的**自我选择接受**机制，直接针对相同的块扩散草稿模型/自回归验证器不匹配问题。链奖励（第5.5节）与streak蒸馏具有相同的核心思想，但在已由\(\mathcal{L}_{\mathrm{CE}}\)物化的教师强制真实词元上评估替代项，仅增加了累积和与指数运算。