Learning to Draft: Adaptive Speculative Decoding with Reinforcement Learning

TL;DR

Learning to Draft is introduced, a novel method that directly optimizes for throughput of each draft-and-verify cycle and is outperforming the state-of-the-art method Eagle3 up to 36.4%.

Abstract

Speculative decoding accelerates large language model (LLM) inference by using a small draft model to generate candidate tokens for a larger target model to verify. The efficacy of this technique hinges on the trade-off between the time spent on drafting candidates and verifying them. However, current state-of-the-art methods rely on a static time allocation, while recent dynamic approaches optimize for proxy metrics like acceptance length, often neglecting the true time cost and treating the drafting and verification phases in isolation. To address these limitations, we introduce Learning to Draft (LTD), a novel method that directly optimizes for throughput of each draft-and-verify cycle. We formulate the problem as a reinforcement learning environment and train two co-adaptive policies to dynamically coordinate the draft and verification phases. This encourages the policies to adapt to each other and explicitly maximize decoding efficiency. We conducted extensive evaluations on five diverse LLMs and four distinct tasks. Our results show that LTD achieves speedup ratios ranging from 2.24x to 4.32x, outperforming the state-of-the-art method Eagle3 up to 36.4%.

Figures (7)

• Figure 1: The overview of LTD method. We formulate the interaction between draft model and target model as an RL environment. We employ two policies to dynamically manage the allocation between draft generation and target verification by controlling the draft tree's depth and verification size. These policies are optimized together using the reward signal based on the throughput of each draft-and-verify cycle.
• Figure 2: Analysis of Draft Time and Verification Time. Left: The average draft and verification latency of Vicuna-13B as a function of draft depth D and verification size V on benchmarks. Right: The composition of the total inference time for Eagle3 and LTD on the benchmarks, with Eagle3's total time normalized to 100% as baselines.
• Figure 3: We report the training reward curves for the Size Policy and Depth Policy of the Vicuna-13B model, as well as the evolution of the speedup ratio on the validation set.
• Figure 4: Left: The effectiveness of our iterative training strategy on Llama3-8B and Dpsk-8B models. Right: The contribution of each component on Vicuna-13B model.
• Figure 5: The average acceptance length for Llama-3.1-8B-Instruct using Eagle3 on HumanEval dataset. Left: Change the draft depth with verification size fixed to 60. Right: Change the verification size with draft depth fixed to 8.