Speculative Sampling with Reinforcement Learning

AI summary

Overview

• Research area: LLM inference acceleration — specifically speculative sampling (SpS) with tree-structured drafting, combined with reinforcement learning for runtime control.
• Technical level: Intermediate. Readers need some familiarity with autoregressive decoding, speculative/draft-verify sampling, and basic reinforcement learning (PPO).
• Scope: This paper introduces Re-SpS, the first reinforcement learning framework that dynamically tunes draft-tree hyperparameters (total tokens, depth, expansion factor) during generation to accelerate LLM inference without changing output fidelity.

What This Paper Is About

State-of-the-art speculative sampling methods such as EAGLE-3 build a "draft tree" of candidate continuations, but the hyperparameters governing that tree — total token budget (TT), depth (d), and top-k expansion factor (k) — are fixed and hand-tuned, so they cannot adapt to different contexts or domains. The paper's goal is to replace those static settings with a learned policy that chooses hyperparameters on the fly, maximizing generation speed while keeping the computational cost of the decision-making itself low. A naive RL implementation would fail because the cost of making decisions (state encoding plus policy inference) can exceed the speedup it produces; the paper's core work is removing that overhead.

Key Contributions

1. Re-SpS, the first RL-based framework for SpS draft tree hyperparameter optimization. The problem of controlling draft tree structure is formulated as a Markov Decision Process (S, A, R, ξ) and solved with online RL (PPO as the backbone).
2. Identification of two concrete sources of RL overhead — state representation overhead (encoding context with an external encoder such as SentenceBERT) and policy inference overhead (a policy forward pass at every decoding step) — and two matching innovations to address them: efficient feature reuse and multi-step action persistence.
3. Extensive empirical evaluation across five benchmarks and three model backbones (LLaMA 3.1-8B, Vicuna-13B, LLaMA 3.3-70B), showing up to 5.45× speedup over the backbone LLM and up to 1.12× over EAGLE-3, with output byte-for-byte identical to greedy decoding.
4. Ablation studies isolating the effect of state representation (text embedding vs. feature vector), cache interval length, and RL algorithm variant (Standard PPO vs. Max-Entropy PPO).

Main Findings

• Headline speedups: Up to 5.45× over the backbone LLM and up to 1.12× over EAGLE-3, with no loss in output fidelity.
• Per-model average gains over EAGLE-3: LLaMA 3.1-8B averages 1.03× (HumanEval 1.07×, Alpaca 1.07×); Vicuna-13B averages 1.04× (HumanEval 1.09×, Alpaca 1.08×); LLaMA 3.3-70B averages 1.06× (HumanEval 1.12×, Alpaca 1.12×).
• Full table (accepted tokens per second, higher is better): For LLaMA 3.1-8B, Re-SpS scores 3.43× (MT-Bench), 3.89× (HumanEval), 3.62× (GSM8K), 3.90× (Alpaca), 2.87× (CNN/DM), mean 3.54×, versus EAGLE-3's mean 3.44×. For Vicuna-13B, Re-SpS: 3.76×, 4.64×, 3.99×, 3.99×, 3.24×, mean 3.92×, versus EAGLE-3's mean 3.80×. For LLaMA 3.3-70B, Re-SpS: 4.47×, 5.45×, 5.13×, 5.34×, 4.03×, mean 4.88×, versus EAGLE-3's mean 4.46×.
• Statistical significance: Differences versus EAGLE-3 on the paired Wilcoxon signed-rank test give p < 10⁻⁴ (LLaMA 3.1-8B), p < 10⁻⁹ (Vicuna-13B), and p < 10⁻²⁹ (LLaMA 3.3-70B). Results in the main table are reported at temperature 0.
• One benchmark regresses: CNN/DailyMail shows slight degradation, reported as 0.98× and 0.97×. The authors attribute this to raising the maximum sequence length from 2048 to 2200 tokens for Re-SpS (to avoid KV cache overflow with long documents) while baselines kept the standard 2048-token limit, adding overhead not present in baseline runs.
• Feature reuse beats text embedding: With Standard PPO, the Feature Vector state gives 1.049× (LLaMA 3.1-8B) vs. 1.044× for Text Embedding, and 1.028× (Vicuna-13B) vs. 1.006×. With Max-Entropy PPO, Feature Vector gives 1.025× (LLaMA 3.1-8B) and 1.033× (Vicuna-13B), versus 1.017× and 1.015× for Text Embedding.
• Max-Entropy PPO produces more diverse behavior: 18 unique actions for LLaMA 3.1-8B and 15 for Vicuna-13B, versus 5 and 3 for Standard PPO — but it is not always the fastest (1.025× vs. 1.049× on LLaMA 3.1-8B; 1.033× vs. 1.028× on Vicuna-13B).
• Longer cache intervals reduce overhead: Increasing the cache interval from 1 to 50 decoding steps substantially reduces inference latency and increases generated tokens per second on LLaMA 3.1-8B. The framework uses N = 10 during training and N = 30 during inference.
• Lossless output: Re-SpS inherits EAGLE-3's lossless property through target-model verification, preserving the output distribution of standard autoregressive decoding; the paper contrasts this with MEDUSA, which relaxes acceptance conditions and lacks rejection-based correction.

Methodology in Plain English

The researchers frame the choice of draft tree hyperparameters as a sequential decision problem. At each decision point the agent observes a state, picks an action, and receives a reward.

• State: Rather than encoding the whole context with an external model (which the authors estimate at roughly 5–15 ms per call and 384-dimensional vectors), the state reuses features the target model already computes — specifically, a concatenation of hidden states from three strategically selected layers, denoted h^(h,m,l)_LM. These are features EAGLE-3 already produces for its draft model, so the state costs no extra inference. EAGLE-3 passes such hidden states through a fully connected layer into one fused vector; Re-SpS instead concatenates the three layers directly, which the authors say reduces computation.
• Action: A tuple (TT, d, k) drawn from finite sets — the upper limit on total tokens, tree depth, and top-k expansion factor per layer. In the implementation, TT ∈ {32, 48, 64, 80, 96, 128}, d ∈ {3, 4, 5, 6, 7, 8}, k ∈ {8, 12, 16, 20, 32}, giving 180 combinations, filtered by the constraint TT ≤ k^(d−1) for computational feasibility.
• Reward: Generation speed measured as accepted tokens divided by elapsed time in seconds. Under action caching, the reward is averaged over the N steps in the cache interval.
• Algorithm: PPO, chosen for stability in sequential decision-making, with a clipped surrogate objective. The paper also tests a maximum-entropy variant that adds an entropy regularization term with weight β_H to encourage exploration. The default configuration uses Max-Entropy PPO with β_H = 0.1.
• Action caching (multi-step action persistence): Instead of querying the policy every decoding step, the selected (TT, d, k) is cached and reused for N consecutive decoding steps, amortizing policy inference cost across many steps. This relies on the Markov property, since the averaged reward already captures the action's temporal impact.
• Training data and setup: A subset of ShareGPT and UltraChat200K with 4,000 questions spanning multiple domains (the combined datasets contain 266,576 questions; Writing is the largest category at 41.9%, followed by General at 21.9% and Reasoning at 20.9%). Policy and value networks are two-layer MLPs with 128 hidden units, implemented via Stable Baselines 3's PPO with [128,128] hidden layers. LLaMA 3.1-8B and Vicuna-13B are trained on a single NVIDIA A40 GPU (stated as 40GB in the main text and 48GB in the appendix), while LLaMA 3.3-70B is trained on four NVIDIA H100 GPUs at 80GB each. Training is one pass over the 4,000 questions; inference is one pass over each benchmark, with each benchmark consisting of 80 questions. Other settings include a learning rate of 3×10⁻⁴, 64 PPO steps, batch size 32, 4 PPO epochs, clip range 0.2, gamma 0.99 (Standard) or 0.95 (Max-Entropy), GAE lambda 0.95 or 0.9, value function coefficient 0.5, and a random seed of 42 for dataset shuffling.
• Evaluation: Five tasks matching EAGLE-3's protocol, using the same weights for all tasks without task-specific fine-tuning: MT-bench (multi-turn conversation), HumanEval (code generation), GSM8K (mathematical reasoning), Alpaca (instruction following), and CNN/DailyMail (summarization).

Why This Matters

Impact on research. The paper reframes speculative sampling configuration from a hand-tuning problem into a sequential decision problem with a learned controller, and it shows that the naive version of that idea is defeated by its own overhead. The two fixes — reusing the target model's internal features as state and caching actions across steps — are general mechanisms that the authors state apply to other tree-based SpS methods including Medusa, EAGLE-2, and EAGLE-3. It also strengthens the case for learned adaptivity over heuristics: prior adaptive work (EAGLE-2's confidence-based pruning, OPT-Tree's fixed-budget node allocation, ProPD's regression-based sizing, bandit methods like MetaSD and BanditSpec, and distillation-based HASS) either keeps structure fixed or selects among pre-defined configurations, rather than learning a context-dependent policy over structural hyperparameters.

Real-world applications (mapped to the benchmark domains the paper evaluates):

• Multi-turn conversational assistants and chat interfaces (MT-bench), where long context and 2–5 turns per question make per-step inference cost dominant.
• Code generation and coding assistants (HumanEval), where Re-SpS shows its largest relative gains over EAGLE-3.
• Mathematical and step-by-step reasoning systems (GSM8K).
• Instruction-following and summarization pipelines (Alpaca and CNN/DailyMail), the latter being the setting where the sequence-length handling issue appears.

Industry relevance. Serving LLMs at scale is latency- and cost-bound, and speculative sampling is already a standard deployment technique. A controller that adds 1.03×–1.12× on top of EAGLE-3 and up to 5.45× over the backbone model translates directly into throughput per GPU. The fact that Re-SpS reuses hidden states already computed by the EAGLE-3 pipeline — rather than adding a separate encoder — is what makes the approach practical rather than a net loss. The released code is at https://github.com/wmd3i/ReSpS.git.

Future Directions

• Extending beyond tree-based SpS. The authors state they will extend the framework to other speculative sampling architectures; the algorithm is described as designed to be efficient and adaptive to all tree-based methods.
• Richer contextual state representations. The paper lists more sophisticated state representations as future work, implying the current three-layer feature concatenation is not the ceiling on context awareness.
• Multi-objective optimization. Future work explicitly targets optimization across throughput, latency, and memory efficiency together, rather than the single speed reward used here.
• Resolving the CNN/DailyMail regression and its evaluation asymmetry. The reported 0.98× and 0.97× came from evaluating Re-SpS at a 2200-token maximum sequence length while baselines used 2048, so the controlled comparison of long-document behavior remains open.
• Open question from the ablations. The paper notes a complex interplay between network capacity, RL algorithm, and target model — Max-Entropy PPO consistently promotes action diversity, but which configuration gives the best speedup depends on the specific model. The additional ablations the paper references in Appendix C are not included in the provided content.

Target Audience

Intermediate-to-advanced machine learning practitioners and systems engineers working on LLM inference serving and latency reduction, especially those already familiar with EAGLE-2/EAGLE-3-style speculative decoding. It also suits reinforcement learning researchers interested in applied sequential control where the agent's own runtime cost is part of the optimization problem, and graduate students looking for a clear example of an RL formulation that fails in its naive form and is rescued by two specific engineering decisions.

Authors’ 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.