PaperlayerView exact revision on arXiv

Source-linked AI summary

Reward-Guided Speculative Decoding for Efficient LLM Reasoning

Baohao Liao, Yuhui Xu, Hanze Dong, Junnan Li, Christof Monz, Silvio Savarese, Doyen Sahoo, Caiming Xiong

arXiv:2501.19324v3cs.CLcs.AI

TL;DR

Large language model inference is costly, and strict unbiased speculative decoding can waste computation on useful draft outputs rejected for distribution mismatch. RSD uses reward-guided mixing and selective target-model refinement to prioritize high-reward steps, achieving improved reasoning accuracy over SD while reducing computation versus target-only decoding.

  • Problem

    Large-scale LLM inference is computationally expensive, while speculative decoding’s strict unbiasedness can reject high-quality draft tokens and limit efficiency on complex reasoning tasks.

  • Method

    RSD dynamically mixes lightweight draft and powerful target models using process rewards to accept high-value draft steps and refine lower-reward steps.

  • Results

    Across reasoning benchmarks, RSD improves accuracy by up to 3.5 over SD on average and reduces FLOPs by up to 4.4× versus using the target model alone.

  • Takeaways & Limitations

    RSD provides a reward-guided inference strategy that balances computational cost and reasoning quality across challenging tasks.

  • Takeaways & Limitations

    The experiments rely on a general process reward model, while specialized or general-domain PRMs are not yet fully available and may require further training or tuning.

Abstract

from arXiv · show

We introduce Reward-Guided Speculative Decoding (RSD), a novel framework aimed at improving the efficiency of inference in large language models (LLMs). RSD synergistically combines a lightweight draft model with a more powerful target model, incorporating a controlled bias to prioritize high-reward outputs, in contrast to existing speculative decoding methods that enforce strict unbiasedness. RSD employs a process reward model to evaluate intermediate decoding steps and dynamically decide whether to invoke the target model, optimizing the trade-off between computational cost and output quality. We theoretically demonstrate that a threshold-based mixture strategy achieves an optimal balance between resource utilization and performance. Extensive evaluations on challenging reasoning benchmarks, including Olympiad-level tasks, show that RSD delivers significant efficiency gains against decoding with the target model only (up to 4.4x fewer FLOPs), while achieving significant better accuracy than parallel decoding method on average (up to +3.5). These results highlight RSD as a robust and cost-effective approach for deploying LLMs in resource-intensive scenarios. The code is available at https://github.com/BaohaoLiao/RSD.

1. Introduction

Large language models offer strong reasoning capabilities but impose substantial inference costs, especially for long, multi-step tasks. RSD addresses speculative decoding’s strict unbiasedness constraint by using reward-guided draft acceptance and target-model refinement to balance efficiency and accuracy.

  • Motivation: Inference at scale creates major computational and economic costs, motivating techniques that reduce energy use and improve deployment scalability.The paper identifies inference, rather than only training, as a key challenge for serving large numbers of queries.
  • Existing approach: Speculative decoding uses a lightweight model to generate candidates that a larger model validates and refines, reducing sequential decoding costs.The approach is intended to reduce latency for real-time or high-throughput applications.
  • Existing approach: Strict unbiasedness can reject high-quality draft tokens whose probabilities under the large model are low, wasting computation during long-horizon reasoning.This limitation restricts flexibility in exploring diverse completions and can negate speculative-decoding speedups when models diverge.
  • RSD: RSD evaluates draft steps with reward signals, accepts sufficiently valuable outputs, and invokes the target model for lower-reward steps instead of discarding all mismatches.The adaptive process mixes lightweight draft evaluation with reward-driven target-model refinement.
  • Results: Across GSM8K, MATH500, Olympiad Bench, GPQA, MMLU STEM, and GaoKao-2023-En, RSD improves reasoning accuracy by up to 3.5 over SD on average and uses up to 4.4× fewer FLOPs than the target model alone.The reported gains target the efficiency–accuracy trade-off in reasoning-intensive inference.

2. Reward-Guided Speculative Decoding

RSD forms a dynamic mixture of draft- and target-model distributions, weighting each according to the reward of the candidate step. High-reward outputs favor the cheaper draft model, while low-reward outputs shift weight toward the target model as a reliable fallback.

  • Setup: At each generation step, the context combines the prompt with previously generated outputs, and either the draft or target model samples the next step.The draft model is denoted m and the target model M.
  • Reward-guided mixture: The reward function evaluates each candidate step’s quality within the response sequence, with higher reward indicating better alignment with the prompt and preceding outputs.The reward is defined conditionally on the current context.
  • Reward-guided mixture: PRSD is defined as a dynamic mixture of Pm and PM whose weights depend on the conditional output’s quality.The target contribution remains present as a stable fallback, while the draft contribution varies with reward.
  • Reward-guided mixture: When reward is high, ω approaches 1 and the draft distribution dominates; when reward is low, the mixture shifts toward the target distribution.This assigns cheaper draft generation to preferred outputs and uses the larger model for lower-quality outputs.
  • Algorithm: RSD dynamically mixes the models at each step so higher-quality steps are more likely to come from the target model while cost-effective draft generation is used when possible.The algorithm’s objective is to balance generation quality and computational efficiency.

2. Compute Reward: The reward function r(yi | zi),

RSD uses reward-guided weighting and rejection sampling to mix draft and target models, favoring high-reward steps while invoking the target selectively. Its theory identifies threshold-based binary weighting as optimal under a sampling budget, and experiments report stronger reward and efficiency than simpler alternatives.

  • RSD procedure: RSD dynamically mixes draft and target outputs, using rejection sampling to replace rejected draft steps with target-generated steps.Accepted steps come from the draft distribution; rejected steps come from the target distribution.
  • RSD procedure: The algorithm uses the draft model when generated steps meet the reward criterion, reducing computation compared with always using the target model.The dynamic weighting function determines when draft steps are of sufficient quality for acceptance.
  • Theoretical guarantees: The induced RSD distribution has expected reward at least as high as the draft model under the stated assumptions.The analysis assumes the target model is at least as strong as the draft model in expected reward at each step.
  • Weighting-function evaluation: Among evaluated weighting functions, binary weighting achieves the best performance and surpasses the single target model, while constant weighting performs worst.All candidates outperform the draft model alone, and reward feedback is important for the weighting decision.
  • Optimal weighting: The optimal constrained-budget strategy is a binary step function that assigns weight 1 to high-reward outputs and weight 0 to lower-reward outputs.The threshold is selected to satisfy the sampling-budget constraint.
  • Reward comparison: Figure 3 compares draft and target reward scores across all questions, correctly answered questions, and reward-based winning rates.For correctly answered MATH500 questions, the target model consistently has higher reward than the draft model.

3. Empirical Results

Across reasoning benchmarks, RSD generally offers a favorable accuracy–efficiency trade-off, outperforming several smaller-model scaling and speculative-decoding baselines. Its threshold controls compute allocation, while reward-guided weighting and target-model corrections improve performance on challenging tasks.

  • Benchmark Setup: RSD is evaluated with general-purpose and math-focused draft and target models across seven reasoning benchmarks.The evaluation includes GSM8K, MATH500, MMLU STEM, OlympiadBench, GaoKao-2023-En, GPQA, and Minerva Math.
  • Benchmark Comparisons: Test-time scaling methods using extensive draft-model sampling consistently underperform a single target model on average.The comparison includes majority voting, Best-of-N, beam search, and process Best-of-N.
  • Benchmark Comparisons: Speculative decoding often underperforms in practice despite theoretically matching the target model’s accuracy through unbiasedness.The paper reports this discrepancy across its reasoning-benchmark evaluation.
  • Search-Based Comparisons: RSD outperforms beam search and process Best-of-N across all three evaluated benchmarks.The authors attribute this pattern to RSD’s larger-model generation and step-wise PRM guidance rather than exhaustive search.
  • Efficiency: RSD (1.5B/7B/7B) improves accuracy by 1.2 and 1.4 over SD (1.5B/7B) and Target (7B), respectively, while using fewer FLOPs.On MATH500, RSD (7B/72B/7B) reaches 88.0 accuracy.
  • Threshold Ablation: Accuracy peaks at δ = 0.7, and once δ exceeds 0.6 it remains higher than using only the target model.At δ = 0.7, the draft model alone solves 48% of questions, avoiding target-model involvement on those questions.
  • Weighting-Function Ablation: The binary step weighting function performs best, surpasses the single target model, and provides δ-based control over inference cost.All weighting candidates outperform the draft model alone, while constant weighting performs worst.

4. Discussion

Discussion experiments examine overhead, model merging, reward-model choice, and general-domain use. The results support RSD’s flexibility, while identifying specialized reward modeling and general-domain reward coverage as open boundaries.

  • Overheads: A 1.5B PRM incurs minimal additional overhead compared with SD, despite being invoked about 18 times per MATH500 question.The reported average is 18 reasoning steps per question, and the tiny PRM still outperforms the single target model in RSD accuracy.
  • Model Merge: Merging models does not necessarily degrade performance and remains superior to SD.Merging larger models can even improve performance while reducing the number of served models.
  • General-Domain Task: General-domain generation lacks an established PRM, so the paper tests an ORM as an alternative reward signal.The authors state that a general-domain PRM and dedicated δ tuning could further improve performance.
  • General-Domain Task: With an ORM on AlpacaEval, RSD achieves a significantly better win rate than the draft model, while 65% of tokens come from the draft model without target intervention.The experiment uses Skywork-Reward-Llama-3.1-8B-v0.2 as the ORM.
  • Combination with SD: RSD can be combined with SD to regenerate rejected steps and further optimize efficiency without additional costs.This combination uses draft-plus-target speculative decoding during rejection events.
  • Limitations: A specialized PRM aligned with the draft model could increase acceptance rates, but training or fine-tuning such a model is left for future work.The current experiments rely on an open-source general PRM.

5. Related Work

Related work spans speculative and parallel decoding, search-based extensions, and reward models for reasoning. RSD differs by using process rewards for stepwise speculative reasoning rather than only token verification or final-output scoring.

  • Speculative Decoding: Speculative decoding uses a lightweight draft model to predict tokens that a larger model verifies in parallel, enabling lossless acceleration.Tree-based speculation extends this approach by generating multiple candidates, while self-speculative decoding reuses parts of the large model.
  • Speculative Decoding: Parallel decoding methods streamline generation with draft models, whereas RSD uses process rewards to perform stepwise speculative reasoning.This distinction is presented as RSD’s difference from previous speculative-decoding methods.
  • Reward Models on Reasoning: Outcome reward models score final outputs, while process reward models provide dense step-level signals based on reasoning-step annotations.These reward models support selecting correct reasoning trajectories during training and inference.

6. Conclusion.

RSD improves LLM inference efficiency for reasoning-intensive tasks by combining draft and target models with reward-guided selective refinement. Evaluations report efficiency gains while maintaining accuracy across reasoning benchmarks.

  • RSD dynamically combines a lightweight draft model with a more capable target model, using process rewards to guide output selection at each step.Outputs are selectively refined based on reward scores.
  • RSD achieves significant efficiency gains over SD and BoN while maintaining accuracy benchmarks.
  • Extensive evaluations across reasoning tasks highlight RSD’s robustness, adaptability, and effectiveness for LLM deployment.

Impact Statement

The paper frames RSD as a contribution toward more efficient and sustainable large-scale AI deployment. It also acknowledges that wider adoption may introduce risks requiring responsible use and safeguards.

  • RSD aims to improve the efficiency and scalability of large language model inference.The stated goal is to advance machine learning through more efficient inference.
  • RSD may facilitate more energy-efficient and cost-effective use of computational resources for deploying large-scale AI systems.
  • Greater accessibility and adoption of efficient language models may exacerbate misuse risks, including misinformation generation and biased outputs.The paper encourages responsible application and safeguards to mitigate these risks.
  • The authors characterize the work’s broader alignment with ethical and sustainable AI development as having no immediate negative societal consequences requiring further discussion.

A.2. Proof of Proposition 2.2

The proof establishes that reward-guided mixing can achieve at least the draft model’s expected reward under stated monotonicity and target-versus-draft expectation conditions. It then characterizes the optimal mixture as threshold-based.

  • A.2. Proof of Proposition 2.2: If ω(r) is non-decreasing in r and E_PM[r(y|z)] ≥ E_Pm[r(y|z)], then E_PRSD[r(y|z)] ≥ E_Pm[r(y|z)].The proof identifies covariance and expectation-difference terms as non-negative under these conditions.
  • A.2. Proof of Proposition 2.2: The threshold solution satisfies primal feasibility, dual feasibility, and stationarity conditions.When the constraint is tight, λ ≥ 0; when slack, λ = 0.
  • A.2. Proof of Proposition 2.2: The optimal sampling strategy uses the draft model for higher-reward outputs and the target model for lower-reward outputs.A threshold is selected to satisfy E_Pm[ω_r(ỹ)] ≥ 1 − γ.

B. Additional Empirical Results

The appendix reports accuracy comparisons and examines the threshold δ used to accept draft reasoning steps. It notes that strict unbiasedness can disadvantage speculative decoding when the draft model is stronger than the target model, while threshold choice affects accuracy.

  • B. Additional Empirical Results: Strict unbiasedness can make SD less accurate than the draft model when the draft model outperforms the target model, unlike RSD.
  • B. Additional Empirical Results: A fixed δ may be suboptimal because reasoning tasks have diverse reward distributions and complexities.The paper identifies threshold selection as an accuracy-sensitive design choice.
  • B. Additional Empirical Results: δ = 0.7 is reported as a reliable choice across various settings, while adjustments among 0.6, 0.7, 0.8, and 0.9 can further improve performance.

B.2. Different Reasoning Complexity

RSD benefits from involving the target model more heavily as question complexity increases. At δ = 0.7, draft-only solving falls from 84% on level 1 questions to 19% on level 5 questions, while δ = 0.7 generally performs well across models and tasks.

  • +16.4 accuracy improvement occurs for level 4 questions when the target model is involved instead of using the draft model alone.
  • +15.7 accuracy improvement occurs for level 5 questions when the target model is involved instead of using the draft model alone.
  • At δ = 0.7, draft-only solving decreases from 84% for level 1 questions to 19% for level 5 questions.The corresponding proportions are 67% for level 2, 58% for level 3, and 44% for level 4.
  • The generally effective δ = 0.7 setting can be slightly tuned because task complexity varies across models and tasks.

B.3. Model Merge

The model-merge study uses linear merging to reduce the number of models required for RSD. Merged models outperform speculative decoding, with larger models showing a greater improvement, while more advanced merging methods remain unexplored.

  • Linear merging combines either the target model with the PRM or the draft model with the PRM to reduce the number of models required by RSD.The study uses MergeKit and focuses on the simplest merging strategy.
  • Advanced merging methods are left for future work beyond the evaluated linear-merging strategy.
  • Merged models outperform speculative decoding overall.
  • +1.4 versus +0.8 improvement is observed for larger versus smaller models, respectively, after merging.The passage reports this comparison as evidence that merging helps larger models more substantially.
Loading 2501.19324v3…