Speculative decoding accelerates memory-bound LLM inference without quality degradation by using a fast drafter to propose multiple candidate tokens and the target model to verify them in parallel. However, conventional sequential speculative decoding suffers from mutual waiting between drafting and verification, and repeated exchange of intermediate states further increases memory access overhead. Parallel speculative decoding addresses this limitation by performing drafting and verification within a single target forward pass, allowing future drafts to be prepared while current candidates are being verified. Although effective at small batch sizes, existing parallel speculative decoding methods either require costly continual pretraining with quality degradation or suffer from low acceptance rates. More importantly, this paradigm inherently suffers from uncertainty in both the bonus token and the accepted length, leading to draft verification mismatch and causing throughput gains to collapse at large batch sizes. To address these limitations, we introduce FlexDraft, a lossless speculative decoding framework that flexibly adapts to varying batch sizes through three key designs. (1) Attention Tuning enables block diffusion drafting by tuning only the attention projectors of the final few layers on mask tokens, while keeping the autoregressive path frozen to preserve the target distribution and produce high quality drafts with minimal trainable parameters. (2) Bonus-guided Calibration uses a lightweight MLP conditioned on the resolved bonus token to calibrate draft logits, mitigating draft verification mismatch caused by bonus token uncertainty. (3) Flex Decoding dynamically switches between parallel draft and verify at small batch sizes and sequential draft then verify at large batch sizes, and adjusts verification length based on draft confidence to eliminate redundant computation.


翻译:投机解码通过使用快速草稿器生成多个候选令牌,并由目标模型并行验证,在保证生成质量无损的前提下加速内存受限的大语言模型推理。然而,传统的串行投机解码存在草稿生成与验证之间的相互等待问题,且中间状态的反复交换进一步增加了内存访问开销。并行投机解码通过在单次目标前向传播中同时完成草稿生成与验证,使得当前候选令牌被验证时即可准备后续草稿,从而克服了这一局限。尽管在小批量尺寸下效果显著,现有并行投机解码方法要么需要代价高昂的持续预训练且伴随质量下降,要么面临低接受率问题。更重要的是,该范式本质上存在奖励令牌与接受长度双重不确定性,导致草稿验证失配,使大批量尺寸下吞吐量增益急剧下降。为解决这些问题,我们提出FlexDraft——一种无损投机解码框架,通过三项关键设计灵活适应不同批量尺寸:(1)注意力调优仅需微调最后几层的注意力投影器作用于掩码令牌,保持自回归路径冻结以维持目标分布,从而以最少可训练参数生成高质量草稿;(2)奖励引导校准利用基于已解析奖励令牌的轻量级MLP校准草稿logits,缓解奖励令牌不确定性导致的草稿验证失配;(3)灵活解码在批量尺寸较小时动态切换为并行草稿与验证模式,较大时切换为串行草稿后验证模式,并根据草稿置信度调整验证长度以消除冗余计算。

0
下载
关闭预览

相关内容

TransMLA:多头潜在注意力(MLA)即为所需
专知会员服务
23+阅读 · 2025年2月13日
一文读懂自注意力机制:8大步骤图解+代码
新智元
153+阅读 · 2019年11月26日
国家自然科学基金
20+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
Arxiv
0+阅读 · 5月19日
VIP会员
最新内容
学习数据的几何:形状空间分析数学综述
专知会员服务
6+阅读 · 6月17日
定向能反无人机系统最新发展动态
专知会员服务
7+阅读 · 6月17日
从燃煤战舰到算法战争:水面指挥的永恒要求
专知会员服务
4+阅读 · 6月17日
相关VIP内容
相关基金
国家自然科学基金
20+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
Top
微信扫码咨询专知VIP会员