提出FlexDraft框架,让大模型推理更快更稳,支持不同批量灵活切换。
FlexDraft: Flexible Speculative Decoding via Attention Tuning and Bonus-Guided Calibration

- 通过调整注意力层实现高效拟稿,参数少且保持生成质量。
- 用轻量校准机制解决拟稿与验证不匹配问题,提升接受率。
- 动态切换推理模式,小批量并行、大批量串行,避免性能崩溃。
推测解码通过快速拟稿模型生成候选词,由目标模型并行验证以加速内存受限的大语言模型推理,且不降低质量。但传统顺序式推测解码存在拟稿与验证相互等待的问题,中间状态反复交换进一步增加内存访问开销。并行推测解码通过单次目标前向传播完成拟稿与验证,可在当前候选验证时准备未来拟稿,解决了该瓶颈。然而现有方法在小批量下有效,却需昂贵持续预训练且质量下降,或接受率低。更重要的是,该范式固有地存在奖励词与接受长度的不确定性,导致拟稿验证不匹配,使大批量下吞吐量大幅下降。为此,我们提出无损的FlexDraft框架,通过三项关键设计灵活适配不同批量:(1)注意力调优通过仅调整最后几层掩码词的注意力投影器,实现块扩散拟稿,同时冻结自回归路径以保持目标分布,生成高质量拟稿且可训练参数极少;(2)奖励引导校准使用轻量MLP基于已解析的奖励词对拟稿逻辑值进行校准,缓解因奖励词不确定性引发的验证不匹配;(3)灵活解码在小批量时采用并行拟稿验证,在大批量时切换为串行模式,并根据拟稿置信度动态调整验证长度,消除冗余计算。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。