RADAR: Accelerate Large Language Model Inference With RL-Based Dynamic Draft Trees

AI summary

Overview

Research area: Efficient large language model (LLM) inference, specifically lossless speculative sampling / speculative decoding and reinforcement learning for dynamic draft-structure control.

Technical level: Advanced. The paper assumes familiarity with speculative sampling, draft-verify pipelines, tree attention, Markov Decision Processes, and policy-gradient reinforcement learning.

Scope: The paper proposes RADAR, an offline-RL method that trains a lightweight LSTM prediction model to decide in real time how many times to call the draft model during speculative decoding, producing dynamic (rather than fixed-depth) draft trees.

What This Paper Is About

In speculative sampling, a small draft model proposes candidate tokens that a large target model verifies in parallel, and the number of calls made to the draft model is normally a fixed preset hyperparameter (for example, 8). This fixed budget wastes computation whenever those draft tokens are rejected. RADAR keeps the guarantee that output quality is unchanged, but replaces the fixed budget with a learned, context-aware controller that decides after every drafting step whether to continue or stop, extracting more speed from the same speculative sampling algorithm.

Key Contributions

1. A speculative sampling framework with dynamic draft trees (RADAR) that uses a lightweight prediction model to adaptively decide the number of calls to the draft model during the draft stage, instead of relying on a preset hyperparameter.

2. An MDP formulation of draft tree generation that sidesteps the lack of labeled data: the paper notes that acceptance length is stochastic because of rejection sampling, so the depth of a draft tree cannot be directly labeled. Modeling drafting as a Markov Decision Process with intrinsic and extrinsic rewards removes that requirement.

3. A new offline dataset of acceptance length distributions. Rewards are computed from acceptance length distributions obtained by running the speculative sampling algorithm on the ShareGPT dataset (the repository references the anon8231489123 ShareGPT_Vicuna_unfiltered dataset). For training, an acceptance length is sampled from the relevant distribution and used to compute the action reward, avoiding expensive real-time interaction with the LLM.

4. Empirical validation across three target LLMs and four tasks, reporting a 3.17x–4.82x speedup over the auto-regressive decoding baseline, a 3%–12.9% improvement over EAGLE-3, and a 9.3%–34.3% reduction in draft-model calls relative to EAGLE-3's fixed 8 calls (18.7% average).

Main Findings

• End-to-end speedup: RADAR achieves approximately 3.17x–4.82x speedup over vanilla auto-regressive decoding, and a 3%–12.9% improvement over EAGLE-3. It attains the highest speedup ratio on most of the tested tasks and target LLMs.

• Acceptance length nearly preserved: RADAR's average acceptance length τ is only about 1.2% lower than EAGLE-3's. The paper notes this deviates from the usual pattern in which a higher acceptance length implies a greater speedup, because RADAR wins by cutting redundant draft calls.

• Fewer draft-model calls: RADAR reduces the average number of calls per draft-verification cycle by 9.3%–34.3% versus EAGLE-3's fixed 8 calls, averaging 18.7%. Reported averages are 5.25/6.19/6.20/6.60 (LLaMA-Instruct 3.1 8B), 6.88/7.26/6.83/7.26 (Vicuna 13B), and 6.10/7.20/5.85/6.47 (DeepSeek-R1-Distill-LLaMA 8B) on MT-bench, GSM8K, Alpaca and MBPP respectively.

• Motivating failure case: Running EAGLE-3's source code on MT-bench with LLaMA-Instruct 3.1 8B, draft tokens are completely rejected at a frequency of about 31%, yet the draft model is still called 8 times.

• Best RADAR results in Table 1 (temperature = 1): LLaMA-Instruct 3.1 8B — 3.41x/τ 4.48 (MT-bench), 4.82x/τ 5.32 (GSM8K), 4.04x/τ 5.51 (Alpaca), 4.44x/τ 6.00 (MBPP). Vicuna 13B — 4.05x/τ 5.67, 4.36x/τ 5.87, 3.84x/τ 5.48, 4.75x/τ 6.42. DeepSeek-R1-Distill-LLaMA 8B — 3.86x/τ 4.85, 4.71x/τ 6.33, 3.17x/τ 4.44, 3.99x/τ 5.31.

• How the draft tree changes: On MT-bench with LLaMA-Instruct 3.1 8B at temperature 1.0, RADAR significantly reduces the probability of an acceptance length of 0 compared with EAGLE-3, raises the frequency of acceptance lengths from 1 to 7, and slightly decreases the frequency at length 8. The mechanism turns the draft tree from tall and thin into shorter and wider, raising acceptance probability at lower depths.

• Policy architecture matters: An LSTM policy beats an MLP policy on all four datasets (e.g., MT-bench 3.41x versus 3.01x; MBPP 4.44x versus 4.15x), supporting the view that drafting is a sequential decision problem. Feeding confidence scores beats feeding the draft model's hidden states (e.g., MT-bench 3.41x versus 3.10x), even though hidden states carry richer context.

• Why confidence scores work: The authors measure a Pearson correlation of r = 0.4719 (p = 7.45e-137) between the top-1 confidence score and acceptance length.

• Reward penalty is robust: Throughput curves across three target models are bell-shaped over α in [0.01, 0.08], with consistently strong performance near α ≈ 0.05; the paper interprets α as a Lagrange multiplier trading terminal throughput against draft computation. Each model's training finishes within 30 minutes.

• Cross-device robustness: With LLaMA-Instruct 3.1 8B on MT-bench, RADAR variants trained with mismatched latency profiles still perform well (e.g., RADAR-5090 reaches 131.82 tokens/sec on RTX5090 versus EAGLE-3's 121.41). Training the prediction model costs only about 60 TFLOPs and 30 minutes, making device-specific tuning affordable.

• Limited to depth control: The current method controls only draft tree depth; per-branch pruning or width adaptation are left open.

Methodology in Plain English

RADAR keeps the standard speculative sampling setup — a draft model proposes tokens, the target model verifies them in parallel — and inserts a third component: a small prediction model that sits between drafting steps. After each drafting step, the draft model returns confidence scores for its top-k candidate tokens; the prediction model (a single-layer LSTM followed by a two-layer MLP head, with input_size=10, hidden_size=128, num_layers=1, dropout=0.1, and an MLP hidden size of 128 with ReLU and dropout 0.1) reads those scores plus its own recurrent hidden state and emits a continue or stop signal. On "continue," the draft model runs another forward pass and the tree grows; on "stop," the tree is finalized, serialized, and verified by the target model.

Because acceptance length is random under rejection sampling, the authors cannot label the "right" stopping depth directly. Instead, they model the drafting loop as a Markov Decision Process: the state is the vector of top-k confidence scores, the action is continue (1) or stop (0), the transition either advances the draft model or moves to an absorbing terminal state, and the reward is a constant penalty −α per step plus a terminal reward equal to acceptance length divided by an estimated total generation time. That time estimate sums a fixed overhead, a per-call draft latency, and the prediction model's own inference cost, with the discount factor set to γ = 0.99.

Training data comes from an offline pipeline: 1,000 prompts are sampled from ShareGPT, EAGLE-3 is run with the number of draft calls enumerated from 1 to t_max, top-k confidence scores are recorded as state sequences, and the distribution of acceptance lengths for each call count is derived analytically from node acceptance and rejection probabilities. During training, the state sequence is fed to the prediction model to produce a trajectory, the number of calls from that trajectory selects a distribution, an acceptance length is sampled from it, and REINFORCE updates the parameters (learning rate 1e-4, batch size 256, 100 epochs). Because the data distribution and the policy's behavior match what online RL would see, the paper argues the extrapolation error that normally plagues offline RL is avoided. Evaluation uses k = 10, a maximum of 8 draft calls, temperature 1.0, batch size 1, and 2x NVIDIA RTX3090 GPUs, with all inference tests run in the same environment.

Why This Matters

Speculative decoding is one of the main practical levers for making LLM serving cheaper, but its fixed draft budget means the same amount of drafting work is spent regardless of how predictable the current text is. RADAR shows that a controller small enough to train in under 30 minutes can decide that budget on the fly, while remaining compatible with state-of-the-art tree-based methods such as EAGLE-3 — something the paper argues prior adaptive methods (SpecDec++, DISCO) cannot do because they were designed for chain-based drafts. The paper also positions RADAR as orthogonal to and composable with asynchronous frameworks like PEARL.

Real-world applications include:

• Interactive chat assistants and multi-turn agents, where MT-bench results (up to 4.05x speedup on Vicuna 13B) translate into lower per-token latency.
• Code generation and completion tools, where MBPP results (up to 4.75x) mean faster suggestions inside editors and CI pipelines.
• Reasoning and mathematical assistants, where GSM8K results with DeepSeek-R1-Distill-LLaMA 8B (4.71x) matter for long chains of generated steps.
• Cost-constrained or on-premise serving, where reducing draft-model calls by an average of 18.7% directly reduces wasted GPU work without changing output distribution.

Industry relevance is direct: RADAR requires no fine-tuning of the target LLM, no change to output quality, and drops into the EAGLE-3 codebase, so it is a low-risk optimization for existing speculative decoding deployments.

Future Directions

• Finer-grained topological control. RADAR currently controls only tree depth; extending it to per-branch pruning or width adaptation within an offline RL framework remains an open challenge.
• Scaling the prediction model and its inputs. The current controller consumes only the top-k confidence scores through a lightweight LSTM; larger models and additional features such as positional context could improve termination decisions.
• Composition with asynchronous execution. The paper states RADAR and PEARL can be integrated to achieve even better acceleration, but reports no combined experiment.
• Hardware-portable reward parameters. Because the reward depends on device-profiled latency terms, cross-device generalization and the α trade-off remain areas for further study, even though mismatched profiles performed well in the reported tests.

Target Audience

This paper is most useful to inference-systems engineers and applied researchers who already work with speculative decoding or serving stacks built on EAGLE-family methods, and to reinforcement learning researchers interested in an offline-RL formulation where reward is derived from an analytically computed distribution rather than labeled data. Readers without background in speculative sampling, tree attention, or policy-gradient methods will need to consult the cited preliminaries first.

Note: the paper does not report the size of the evaluation datasets (MT-bench, GSM8K, Alpaca, MBPP), and it explicitly does not evaluate generation quality, stating that its use of a strict speculative sampling algorithm preserves performance.

Authors’ abstract

Inference with modern Large Language Models (LLMs) is expensive and slow, and speculative sampling has emerged as an effective solution to this problem. However, the number of calls to the draft model for generating candidate tokens in speculative sampling is a preset hyperparameter, lacking flexibility. To generate and utilize the candidate tokens more effectively, we propose RADAR, a novel speculative sampling method with RL-based dynamic draft trees. RADAR formulates the draft tree generation process as a Markov Decision Process (MDP) and employs offline reinforcement learning to train a prediction model, which enables real-time decisions on calls to the draft model, reducing redundant computations and further accelerating inference. Evaluations across three LLMs and four tasks show that RADAR achieves a speedup of 3.17x-4.82x over the auto-regressive decoding baseline. The code is available at https://github.com/minaduki-sora/RADAR.