Speculative Sampling with Reinforcement Learning

TL;DR

This work tackles the latency of large language model inference by learning to adapt speculative sampling draft-tree hyperparameters with reinforcement learning. The proposed Re-SpS framework uses efficient state representations drawn from target-model hidden states and a multi-step action persistence mechanism to amortize policy costs, achieving up to 5.45× speedup over the backbone LLM and up to 1.12× over SOTA EAGLE-3 while preserving exact output fidelity. Through evaluations on five benchmarks across multiple model sizes and ablations, Re-SpS demonstrates robust, context-aware dynamic adaptation that outperforms static configurations. The findings indicate that learning-driven adaptive draft-tree control can meaningfully enhance inference throughput in real-world LLM deployments without sacrificing correctness or quality.

Abstract

Inference time latency has remained an open challenge for real world applications of large language models (LLMs). State-of-the-art (SOTA) speculative sampling (SpS) methods for LLMs, like EAGLE-3, use tree-based drafting to explore multiple candidate continuations in parallel. However, the hyperparameters controlling the tree structure are static, which limits flexibility and efficiency across diverse contexts and domains. We introduce Reinforcement learning for Speculative Sampling (Re-SpS), the first reinforcement learning (RL)-based framework for draft tree hyperparameter optimization. Re-SpS dynamically adjusts draft tree hyperparameters in real-time, learning context-aware policies that maximize generation speed by balancing speculative aggression with computational overhead. It leverages efficient state representations from target model hidden states and introduces multi-step action persistence for better context modeling. Evaluation results across five diverse benchmarks demonstrate consistent improvements over the SOTA method EAGLE-3, achieving up to 5.45$\times$ speedup over the backbone LLM and up to 1.12$\times$ speedup compared to EAGLE-3 across five diverse benchmarks, with no loss in output fidelity.

Figures (3)

• Figure 1: Re-SpS vs. EAGLE-2 & 3: Comparison of speculative sampling tree structures through steps $1\cdots n$ of a generation task. Top: EAGLE-2 & 3 uses static hyperparameters (upper limits for the total token $TT$, depth $d$, and expansion factor $k$). Bottom: Re-SpS dynamically adapts the draft tree hyperparameters based on the context. Step $t+1$ exemplifies a more cautious draft tree, tuning top-k up and depth down to safely capture more possible tokens without risking costly rejections. Step $t+n$ shows a more aggressive configuration, with depth adjusted to allow for confident drafting. Draft tokens and branching structures were manually chosen to provide a readable, illustrative comparison. Actual model outputs may vary, but the figure faithfully reflects the parameter and adaptivity constraints of each approach.
• Figure 2: Architecture of Re-SpS: The diagram illustrates the Re-SpS framework for SpS in LLMs. For each new input prefix (from question tasks $q_1, q_2, q_3, \ldots$), hidden state vector ($h, m, l$) from the target model ($M_t$) is aggregated into $s_t$ and passed to a reinforcement learning (RL) policy. The RL agent outputs draft tree hyperparameters $(TT_{t+1}, d_{t+1}, k_{t+1})$ for the next generation step. The draft model ($M_d$) constructs tree-structured speculative candidates ($DT_1, DT_2, DT_3$), which are verified by the target model. The number of accepted tokens and elapsed time are recorded, and the generation speed $r_{t+1} = \frac{\text{accepted tokens}}{\text{elapsed time}}$ is used as the RL reward. Solid arrows show EAGLE-3's static pipeline; dashed arrows highlight Re-SpS's adaptive, RL-driven control. Note: For clarity, the figure is simplified; in practice, the draft trees $DT$ are redefined multiple times for each new prefix within a question task until a maximum sequence length or end-of-sequence is reached.
• Figure 3: Cache interval length vs. inference latency in seconds and generated speed in tokens per second for LLaMA 3.1-8B. The cache interval length is the number of decoding steps over which the RL policy's action is cached and reused. Tested with a fresh SpS RL policy on random 80 questions from the training dataset.