编辑推荐:
摘要大语言模型的自回归解码特性带来了显著的推理延迟问题,因此推测解码被视为一种极具前景的加速范式。最近,基于扩散的草稿模型(如DFlash)被提出,它们能够在单次前向传播中并行生成草稿令牌,实现了显著的加速效果。然而,DFlash假设每个块内的令牌之间相互条件独立,忽略了相邻
摘要
大语言模型的自回归解码特性带来了显著的推理延迟问题,因此推测解码被视为一种极具前景的加速范式。最近,基于扩散的草稿模型(如DFlash)被提出,它们能够在单次前向传播中并行生成草稿令牌,实现了显著的加速效果。然而,DFlash假设每个块内的令牌之间相互条件独立,忽略了相邻令牌之间固有的顺序依赖关系和句法约束。这种假设导致了不连贯的草稿序列、不准确的语义预测,以及在验证阶段较高的令牌拒绝率,从而限制了整体的加速效率。为解决这一局限,我们提出了一种块级窥视策略,该策略引入了一种逐层残差窥视机制用于训练基于扩散的草稿模型。具体而言,我们的方法在草稿模型的中间层逐步揭示真实令牌的嵌入表示,其中揭示比例由层深度和逐层预测准确度共同决定,使模型能够在保持并行推理的同时捕获块内令牌依赖关系。在Qwen3模型上的大量实验(涵盖三个基准任务)表明,我们的方法相比现有的扩散推测解码实现了显著的加速提升。
引言
大语言模型在数学、编码和信息检索等多个领域展现了卓越的能力[2]、[13]、[24]。然而,由于自回归特性所带来的推理延迟问题,仍然是实时应用中持续存在的挑战[29]。自回归解码范式要求模型逐个生成令牌,每个预测的令牌都需要目标模型进行一次完整的前向传播[38]。随着大语言模型的参数规模呈指数级增长,以追求更强的上下文理解、逻辑推理和生成泛化能力,推理的计算开销和内存占用急剧上升,进一步加剧了延迟瓶颈,并推动了对大语言模型推理加速的广泛研究[25]。
为缓解推理延迟问题,现有的主流优化策略大致可分为模型压缩、解码加速和硬件感知优化。模型压缩方法,包括量化[18]、剪枝[26]、[42]和知识蒸馏[15],以减少模型参数和计算复杂度为代价,部分牺牲了模型精度,但不可避免地会导致性能下降,尤其在复杂推理任务上。为实现无损加速,推测解码被提出[19]。该范式利用轻量级草稿模型生成未来令牌,然后使用大型目标模型并行验证。虽然推测解码实现了无损加速,但它只是将自回归解码从目标模型转移到了草稿模型,并未改变自回归生成的本质[5]。
近年来,得益于掩码语言模型[10]和扩散模型[16]的发展,扩散语言模型能够实现多个令牌的并行生成,解耦了令牌之间的依赖关系,从而显著降低了推理延迟[27]。尽管扩散语言模型与自回归模型之间存在性能差距,但扩散语言模型可以进一步加速草稿过程。基于这一理论,DFlash采用轻量级扩散草稿模型来预测未来令牌[5]。
尽管DFlash取得了成功,但它假设块内的令牌是条件独立的,并在单次前向传播中预测所有被掩码的位置。我们认为,这种假设阻碍了模型学习块内令牌之间的上下文依赖关系。在文本生成场景中,连续语义块内的令牌具有固有的顺序相关性和逻辑约束,后续令牌的生成在很大程度上依赖于同一块中前面令牌的语义和句法信息。DFlash对所有被掩码令牌的并行预测完全忽略了这种上下文依赖关系,导致草稿令牌序列不连贯、语义预测不准确,以及在目标模型验证阶段出现较高的令牌拒绝率。这些缺陷严重限制了基于扩散的草稿模型的质量,削弱了推测解码框架的整体加速效率,并阻碍了扩散辅助的大语言模型推理加速的实用部署。
为了解决上述局限,我们提出了一种新颖的块级窥视策略用于DFlash的训练,以捕获块内令牌之间的上下文依赖关系。具体而言,我们引入了一种逐层残差窥视机制,在草稿模型的中间层逐步揭示真实令牌的嵌入表示。在每一层中,揭示比例由层深度和当前预测准确度共同决定,使得更深且预测准确率更低的层获得更多的上下文指导。随着草稿模型在训练过程中预测准确度的提高,真实信息的暴露逐步减少,最终使草稿模型能够独立预测所有被掩码的位置。同时,我们提出的方法仅修改DFlash的训练过程,仍能在单次前向传播中预测未来草稿令牌。由于窥视操作在推理时被禁用,该机制本身相对于DFlash推理流程不引入任何额外的推理时计算。
为评估我们提出的方法的性能,我们在标准基准任务上进行了一系列实验,包括多轮对话(MT-Bench [44])、数学(GSM8K [9])和编码(HumanEval [6])。实验表明,我们提出的方法能够显著提升推测解码的加速效果。与基线DFlash相比,我们提出的方法实现了最多30%的加速提升,证明了块级窥视策略在捕获令牌块内上下文依赖关系方面的有效性。
各节摘要
推测解码
推测解码[19]是一种利用轻量级草稿模型生成未来令牌,然后使用大型目标模型并行验证的范式。目前,关于推测解码的研究大致可分为两个方向:独立起草和自我起草。[40]首先提出了一个独立起草框架,利用轻量级Transformer进行起草。尽管该方法实现了显著的加速,但它需要从...
预备知识
为了便于理解我们提出的方法,本节介绍DFlash和窥视Transformer的关键组件。我们首先介绍块扩散草稿模型的训练,然后描述窥视Transformer的训练策略。
方法
在本工作中,为了捕获块内令牌之间的上下文依赖关系,我们为基于扩散的草稿模型提出了一种块级窥视策略。此外,为避免信息泄露,我们为块扩散模型设计了一个窥视掩码。我们模型方法的详细图示见Fig. 1,本节剩余部分详细描述该方法。
实验设置
模型与评估 为评估我们提出的方法的性能,我们在Qwen3(1.7B、4B)预训练模型上进行了一系列实验。我们评估了三个类别的任务,包括数学(GSM8K)、多轮对话(MT-Bench)和编码(HumanEval)。对于每个任务,我们使用平均接受长度和相对于自回归目标模型的端到端加速来评估我们提出的方法的性能。
数据集 为了公平比较,我们训练我们的模型...
结论
在本工作中,我们为基于扩散的推测解码提出了一种新颖的训练策略,解决了现有扩散草稿模型中条件独立假设的局限性。通过引入逐层残差窥视机制,我们提出的方法在草稿模型的中间层逐步揭示真实令牌的嵌入表示,其中揭示比例由层深度和逐层预测准确度共同决定,在...之间取得平衡。
CRediT作者贡献声明
Shuheng Wang: 撰写—初稿,监督,方法论,资金获取,形式分析,数据整理,概念化。Dongbai Li: 撰写—审阅与编辑,可视化,数据整理。Shumin Shi: 撰写—审阅与编辑,监督,项目管理,资金获取。Zhen Shen: 撰写—审阅与编辑,验证,资金获取。Hanyao Huang: 撰写—审阅与编辑。
关于生成式AI和AI辅助技术在写作过程中的声明
在本工作中,我们使用了GLM-5.2来润色论文,所有润色内容均经作者审阅和批准。
利益冲突声明
作者声明,不存在任何已知的竞争性财务利益或个人关系可能影响本文所报告的工作。
致谢
我们非常感谢审稿人的辛勤工作。本工作得到了国家自然科学基金(项目编号:61732005, 61671064)、河南省自然科学基金(项目编号:252300423330, 262300421803)、河南省科技攻关计划(项目编号:262102210185)以及南阳理工学院博士科研启动基金(项目编号:NGBJ-2024-12)的支持。
Shuheng Wang 于2023年获得南京理工大学博士学位。目前为南阳理工学院讲师。研究方向涵盖自然语言处理、机器翻译、大语言模型和推测解码。
Shuheng Wang | Dongbai Li | Shumin Shi | Zhen Shen | Hanyao Huang