论文

FlexDraft:通过注意力调整和奖励引导校准实现灵活的 推测解码

FlexDraft: Flexible Speculative Decoding via Attention Tuning and Bonus-Guided Calibration

模型推理投机采样

摘要

推测解码 通过使用快速起草者提出多个候选词元和目标模型来并行验证它们,从而加速内存绑定的 LLM 推理,而不会降低质量。然而,传统的顺序推测解码在起草和验证之间存在相互等待的问题,并且中间状态的重复交换进一步增加了内存访问开销。并行 推测解码 通过在单个目标前向传递中执行起草和验证来解决此限制,从而允许在验证当前候选者的同时准备未来的草案。尽管在小批量下有效,但现有的并行 推测解码 方法要么需要昂贵的持续预训练,质量会下降,要么接受率较低。更重要的是,这种范式本质上受到奖励词元和接受长度的不确定性的影响,导致草稿验证不匹配并导致吞吐量增益在大批量下崩溃。为了解决这些限制,我们引入了 FlexDraft,这是一种无损 推测解码 框架,它通过三个关键设计灵活适应不同的批量大小。 (1) 注意力调整通过仅调整掩模标记上最后几层的注意力投影仪来实现块扩散绘图,同时保持自回归路径冻结以保留目标分布并以最小的可训练参数生成高质量的草图。 (2) 额外词元引导校准使用以已解析额外词元为条件的轻量级 MLP 来校准草稿 logits,从而减轻由于额外词元不确定性而导致的草稿验证不匹配。 (3) Flex Decoding动态切换小批量并行草稿验证和大批量连续草稿验证,并根据草稿置信度调整验证长度以消除冗余计算。