Reward-Guided Speculative Decoding for Efficient LLM Reasoning

$~$

Supplementary Material: https://anonymous.4open.science/r/Rebuttal-supp-6595-C4A2/appendix_rebuttal_ICML.pdf

$~$

---

Q1. Why RSD outperforms the large model?

Please refer to our rebuttal to Reviewer p3vx Q5.

$~$

Q2. Alternative reward models or potential biases

As we mentioned, one might consider alternative formulations:

• Likelihood-Ratio-Based: One could incorporate a term like $\frac{\mathbf{P}_M(y|z)}{\mathbf{P}_m(y|z)}$ to define the weighting function (e.g., $w(y|z)=\min(1, \alpha\,\frac{\mathbf{P}_M(y|z)}{\mathbf{P}_m(y|z)})$). Such a formulation may capture discrepancies between the two models’ distributions, but it may also introduce bias if the likelihood estimates are themselves noisy or miscalibrated.
• Hybrid: One might combine the process reward and likelihood ratio—e.g., $w(y|z)=\min(1, \beta\,r(y|z)\,\frac{\mathbf{P}_M(y|z)}{\mathbf{P}_m(y|z)})$. While this hybrid approach can potentially leverage complementary strengths, its sensitivity to misestimations in either component may induce systematic biases. For example, if the reward function overestimates quality in regions where the draft model already performs well, the weighting may overfavor the draft model, reducing the benefits of the target model’s corrections.

In both cases, under our assumption, the underlying distribution will be better than small model induced one. Whereas, the choice of reward function can significantly affect the acceptance probability, and thus the final mixture distribution. When we choose more weights towards likelihood, it will biased to exact-match of large model distribution, so that the cost would be higher (similar to SD).

You may also consider imperfect reward (Q5).

$~$

Q3. Performance by question difficulty.

We show it in Tab C.4 (in link). One can observe:

• For simpler questions (level=1), RSD and SD performs the same;
• For harder questions (level>2), RSD consistently outperforms SD. With the help of PRM, RSD selectively includes the target model for some reasoning steps, correcting the wrong reasoning step and leading to better performance.

$~$

Q4. Different reward models & variations in reasoning complexity.

In Tab C.5 (in link), we includes more PRMs, Qwen2.5-Math-PRM-7B and Qwen2.5-Math-PRM-72B, released during the reviewing cycle.

• RSD is robust to PRM. Although Skywork-o1-Open-PRM (regression model) and Qwen2.5-Math-PRM (classification model with num_labels=2) are trained differently, they all perform quite well across different tasks.
• $\delta=0.7$ consistently yields strong results across all four PRMs, appearing to be a sweet spot for balancing efficiency and performance when the reward score lies in [0, 1].

We also draw the accuracy and efficiency wrt $\delta$ in Fig C.3 (in link). These two PRMs behave similarly:

• Involving the target model ($\delta \neq 0$) consistently improves accuracy over using the draft model alone ($\delta = 0$), with greater gains on harder questions.
• For the same $\delta$, the proportion of questions solved by the draft model alone decreases with an increasing level, showing that harder questions need more involvement of the target model.

$~$

Q5. Imperfect reward models

In practice, $\hat{r}(y|z)$ is an approximation of the oracle reward $r_{\text{oracle}}(y|z)$. We can derive non-asymptotic convergence bounds under assumptions on the reward model’s error distribution.

Assumption (Sub-Gaussian Error): Assume that for each reasoning step, the estimation error

$$\epsilon(y|z) = \hat{r}(y|z) - r_{\text{oracle}}(y|z)$$

is sub-Gaussian with parameter $\sigma^2$; that is, for any $t > 0$,

$$\Pr\left(|\epsilon(y|z)| \geq t\right) \leq 2\exp\left(-\frac{t^2}{2\sigma^2}\right).$$

Under this assumption, when using $n$ independent observations, standard concentration inequalities imply that, with probability at least $1-\delta$,

$$\left|\frac{1}{n}\sum_{i=1}^n \hat{r}(y_i|z) - \mathbf{E}[r_{\text{oracle}}(y|z)]\right| \leq \sigma \sqrt{\frac{2\log(2/\delta)}{n}}.$$

Even if the reward model is imperfect, the empirical average reward converges to the oracle at $O\Bigl(\sqrt{\frac{\log(1/\delta)}{n}}\Bigr)$. The bias introduced in the acceptance decision and in $\mathbf{P}_{\text{RSD}}$ can be controlled in a non-asymptotic manner. Even with imperfect reward estimates, the impact on the RSD diminishes as more data become available.

$~$

Q6. Open-ended generation tasks.

Refer to Q2 of Reviewer GXeJ.

$~$

Q7. Choice of $\delta$ based on the difficulty.

From Fig C.1, $\delta=0.7$ performs consistently well across difficulty levels, achieving top accuracy in 4/5 levels. Q4 further confirms $\delta=0.7$ works well for different PRMs.

That said, dynamically selecting $\delta$ based on question difficulty (eg using a LLM to predict difficulty and adjust $\delta$ accordingly) could improve performance and efficiency. We leave this for future work.

Q1. The reliance on a process reward model, which may introduce additional overhead.

This is indeed a very practical concern, since RSD utilizes one more model, i.e. the process reward model (PRM), than speculative decoding (SD). However, accoding to our experiments, the overhead is minor from three perspectives:

• The PRM is small: According to Table 2, a 1.5B PRM could already offers consistent better performance than all baselines with different draft and target models.
• The inference cost of PRM is minimal: RSD applies PRM to score the step instead of tokens from the draft model. The average number of steps per question in MATH500 is 18. The inference cost of PRM is similar to generate 18 tokens per question.
• The PRM can be merged with the target or draft model: In Table 4, we explore the possibility of merging the PRM with the draft or target model. Surprisingly, model merging doesn't obviously degrade the performance, and even results in better performance when merging the larger models. In this way, RSD shares the same number of models as SD. A further investigation of model merging between PRM and proxy model is left to future work.

$~$

Q2. How does RSD perform on non-reasoning tasks like text generation or summarization?

This is a very interesting suggestion. An important component of RSD is PRM. As far as we know, there is not yet a PRM for general-domain generation. But there are plenty of outcome reward models (ORMs) for open-ended generation. Could we use ORM instead of PRM in RSD?

Experimental setup: We utilize Llama-3.2-1B-Instruct as the draft model, Llama-3.1-8B-Instruct as the target model, and Skywork-Reward-Llama-3.1-8B-v0.2 as the ORM. The 805 prompts from AlpacaEval [1] are used for the generation, and the model outputs are evaluated with AlpacaEval2.0 against the outputs from gpt4_turbo. Similar to the setting in the paper, we define a generation ended with "\n\n" as a reasoning step, and apply the ORM to score this step. The score of Skywork-Reward-Llama-3.1-8B-v0.2 ranges from -$\infty$ to $\infty$. We didn't extensively tune the reward threshold $\delta$, and empirically chose $\delta=0$.

Results: As shown in the following table, even with an ORM instead of a PRM, RSD achieves a significantly better win rate than the draft model, showing RSD's robustness across different tasks. Among all generated tokens, 65% tokens are generated by the draft model only without any intervention of the target model. We believe that a general-domain PRM and dedicated tuning of $\delta$ could further boost the performance.

$~$

Q3. Could the framework be extended to handle multimodal inputs?

Yes, we believe our RSD could be seamlessly extended to the multomodal reasoning tasks, because:

• The key component of RSD is PRM. As shown in Q2, even an ORM could act as a PRM. We believe that the existing multimodal ORM can be used in RSD for multimodal reasoning tasks.
• We noticed that a new multimodal PRM [2] is released recently. It can be used perfectly in our framework.

However, the exploration of multimodal reasoning tasks is out of the scope of this work, and there isn't any multimodal PRM when we conducted this research, we leave this investigation to interested readers.

$~$

[1] Length-Controlled AlpacaEval: A Simple Way to Debias Automatic Evaluators. Yann Dubois, Balázs Galambosi, Percy Liang, Tatsunori B. Hashimoto

[2] VisualPRM: An Effective Process Reward Model for Multimodal Reasoning. Weiyun Wang, Zhangwei Gao, et. al., Wenhai Wang

$~$

---

Thank you very much for your thoughtful suggestions, making our work more solid. We have incorporated the new results in the updated version by our side.

If these revisions address your concerns, we kindly request a reconsideration of the scores. Should you have any further questions, we are happy to assist.

Q1. PRM vs large models’ likelihood

The core issue is the misalignment between the large model (LM) and PRM, leading SD to reject high-quality tokens due to low LM likelihood. While the reviewer suggests this implies high FP rates from LM, these are often style-/format-related, not correctness errors. LMs generalize better but are NOT trained for correctness. In contrast, PRMs are explicitly trained to identify correctness, making them more suitable for detecting errors, while LM likelihood would point out both style inconsistency and errors. Thus, our algorithm uses LM guidance only after PRM rejection, leveraging both PRM's specialized correctness identification and LM’s strength in generalizing better completions. Using both for rejection is also promising; we leave it for future work as noted in Sec 2.4.

$~$

Q2. Calculation of FLOPs.

As described in Sec 3.3, we follow the standard FLOPs approximation for transformer models with N parameters, i.e., approximately 2N FLOPs per inference token, as adopted in prior works [1,2]. We provide a detailed example below: The FLOPs calculation for RSD (7B/72B/7B) is presented in the table below.

[1] Scaling laws for neural language models.

[2] Beyond chinchilla-optimal: Accounting for inference in language model scaling laws.

$~$

Q3. More efficiency metrics (throughput).

We measure the througput and provide the following table. We use batch size 256 and MATH500. The observation is similar to Figure 4: RSD is faster than SD.

$~$

Q4. Search-based baselines

We compare RSD to the search-based baselines (Process Best-of-N and Beam Search) in Table 3, where RSD significantly outperforms the search-based baselines that only utilize a draft model and PRM. It shows the importance of involving a larger model in the reasoning path for correction.

$~$

Q5. Why RSD outperforms the large model?

Proposition (Improved Expected Reward via RSD): Assume that for each decoding step the reward function $r(y|z)\in[0,1]$ and consider $\omega(r) = \mathbf{1}(r \geq \delta)$ for some threshold $\delta \in [0,1]$. Define the RSD mixture distribution as $\mathbf{P} _ {\text{RSD}}(y| z) = \omega(r(y| z)) \mathbf{P} _ m(y | z) + \nu\mathbf{P} _ M(y| z),$ where $\nu = 1 - \mathbf{P} _ {y\sim \mathbf{P} _ m}\{r(y| z) \geq \delta\}$. The expected reward of the RSD distribution is $\mathbf{E} _ {\mathbf{P} _ {\text{RSD}}}[r(y|z)] = \alpha\,\mathbb{E} _ {\mathbf{P}_m}[r(y| z)| r(y| z) \geq \delta] + (1-\alpha)\,\mathbb{E} _ {\mathbf{P}_M}[r(y| z)],$ with $\alpha = \mathbf{P} _ {y\sim \mathbf{P}_m}\{r(y| z) \geq \delta\}.$

The RSD distribution outperforms the target model $\mathbf{P} _ M$ (i.e., $\mathbf{E} _ {\mathbf{P} _ {\text{RSD}}}[r(y| z)] > \mathbf{E} _ {\mathbf{P} _ M}[r(y| z)]$) if and only if $\mathbf{E} _ {\mathbf{P}_m}[ r(y| z) | r(y| z) \geq \delta] > \mathbf{E} _ {P_M}[r(y| z)].$

Thus, RSD yields a higher expected reward than $\mathbf{P}_M$ provided that the subset of tokens generated by the draft model $\mathbf{P}_m$ (i.e. those with $r(y| z) \geq \delta$) has a higher reward than $\mathbf{P}_M$. If one can choose $\delta$ so that only high-quality tokens from the draft model are accepted, then the mixture of these with the fallback tokens from $\mathbf{P}_M$ lifts the overall expected reward of RSD above that of the large model alone.

This result offers a theoretical explanation for why RSD can outperform the large model: it leverages the fact that, when the draft model’s top-performing outputs (as measured by $r$) are even better on average than the target model’s outputs, selectively accepting these outputs increases the overall quality.

In summary, for RSD to be advantageous over $\mathbf{P}_M$, one must choose a threshold $\delta$ such that: $\mathbf{E} _ {\mathbf{P}_m}[r(y| z)| r(y| z) \geq \delta] > \mathbf{E} _ {\mathbf{P}_M}[r(y| z)].$

Under this condition, RSD leads to an improved expected reward compared to $\mathbf{P}_M$.

$~$

Q6. Estimation of δ

A6. As noted in Prop. 2.3, $\delta$ corresponds to a quantile of the reward distribution under a given compute budget. While theoretical guidance informs its range, estimating the exact quantile requires access to empirical reward statistics -- hence, practical tuning (e.g., grid search) is necessary and suggested by our analysis. Moreover, Table 2 demonstrates that performance is not highly sensitive to $\delta$, suggesting our method is robust across a reasonable range and not reliant on fine-tuning.