FlexDraft: Flexible Speculative Decoding via Attention Fine-tuning and Reward-Guided Calibration

FlexDraft is a lossless speculative decoding framework. To address the performance collapse issue of traditional speculative decoding methods in large-batch scenarios, it achieves flexible adaptation to varying batch sizes through three key designs: attention fine-tuning, reward token-guided calibration, and dynamic decoding strategy switching, thereby improving LLM inference efficiency without sacrificing output quality.

In LLM inference acceleration, speculative decoding amortizes computational costs by having a draft model generate candidate tokens which are then verified in parallel by the target model. However, traditional sequential speculative decoding faces bottlenecks such as mutual waiting between draft generation and verification, and increased memory access overhead. While parallel speculative decoding attempts to solve this problem, existing methods either require expensive pre-training with quality degradation or have low acceptance rates. Moreover, the uncertainty of reward tokens and acceptance lengths leads to a sharp collapse in throughput gains in large-batch scenarios.

FlexDraft adopts an attention fine-tuning strategy: it only fine-tunes the attention projection layers in the last few layers of the target model, trains only on masked tokens, and freezes the autoregressive path. This design preserves the original distribution characteristics of the target model, endows it with the ability to generate high-quality drafts, has low training costs, and the block-level diffusion draft method balances efficiency and effectiveness.

To address the draft-verification mismatch problem caused by the uncertainty of reward tokens in parallel speculative decoding, FlexDraft introduces a lightweight MLP calibration network. It calibrates the draft logits conditional on the resolved reward tokens, effectively alleviating the mismatch problem, improving acceptance rates without significantly increasing inference overhead.

FlexDraft's dynamic strategy switching mechanism automatically selects the optimal decoding strategy based on the current batch size: in small-batch scenarios, it uses the parallel draft-verification mode to maximize throughput; in large-batch scenarios, it switches to the sequential draft-verification mode to avoid performance collapse. It also dynamically adjusts the verification length based on draft confidence to eliminate redundant computations, ensuring efficient inference under different loads.

Compared to model compression techniques like quantization and pruning, FlexDraft is completely lossless (the output distribution is consistent with the original model). Compared to other speculative decoding methods, it has better stability in large-batch scenarios (achieved through reward-guided calibration and dynamic strategy switching). Similar to speculative execution in the CPU domain, it represents an attempt at intelligent scheduling of computing resources in the AI inference field.

FlexDraft demonstrates that an elegant architectural design can achieve efficient lossless speculative decoding. The attention fine-tuning strategy provides a new idea for model adaptation (adjusting only key components without full fine-tuning). The dynamic switching mechanism adapts to dynamic loads in production environments, which is of great significance for building high-throughput, low-latency inference services.

The FlexDraft framework is extensible: in the future, more complex calibration network designs can be explored, or it can be applied to other generation tasks. The dynamic strategy switching mechanism can inspire the design of other adaptive systems. With the rise of multimodal models and agent systems, such efficient inference work will provide important technical accumulation for AI infrastructure.