Reward-Shifted Speculative Sampling

• The paper introduces SSS, a method where an aligned draft model is used to recover an RLHF-optimal reward-shifted distribution without modifying the large target model.
• SSS modifies standard speculative sampling by comparing draft proposals against an SFT baseline and applying bonus-token corrections to ensure proper preference alignment.
• Experimental results on models such as OPT and Llama show that SSS achieves higher reward metrics and improved latency compared to traditional speculative decoders.

Reward-Shifted Speculative Sampling (SSS) is a decoding algorithm for test-time weak-to-strong alignment in which a small draft model is aligned with human preferences, the large target model remains unchanged, and speculative sampling is modified so that generation follows a reward-shifted, RLHF-optimal distribution rather than the target model’s original distribution (Li et al., 20 Aug 2025). In this construction, the aligned draft model supplies the reward shift, the unchanged target model supplies large-model likelihoods, and the acceptance rule plus bonus-token distribution are redesigned so that preference alignment is performed at inference time without post-training the large target model and without querying an external reward model during decoding.

1. Formal objective and alignment setting

SSS is formulated in a KL-regularized RLHF setting. The underlying objective is

$$\max_{\bm{\theta}} \mathbb{E}_{x\in\mathcal{D}_p,\, y\sim\pi_{\bm{\theta}}(\cdot|x)} r(x,y) - \lambda \cdot \mathbf{KL}(\pi_{\bm{\theta}} \| \pi_{\text{ref}}),$$

whose optimal solution is

$$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$

SSS takes this RLHF-optimal policy as its target distribution rather than the base target model’s native autoregressive law (Li et al., 20 Aug 2025).

The method is explicitly framed as test-time weak-to-strong alignment. The weak model is a small draft model, denoted $\pi_{\text{draft}^r}$, which is first SFT-trained on chosen preference responses and then aligned with DPO. The strong model is a larger target model, denoted $\pi_{\text{ref}}$, which remains unchanged. The central claim is that the distributional shift between the aligned draft and the unaligned target can be exploited to recover the RLHF-optimal aligned distribution of the strong model without directly aligning that strong model (Li et al., 20 Aug 2025).

A crucial modeling assumption is that the aligned draft approximately follows the same reward tilt relative to its SFT checkpoint: $$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$ In the theorem, this approximation is strengthened to equality. This assumption is what makes it possible to infer the missing reward factor from the gap between $\pi_{\text{draft}^r}$, $\pi_{\text{draft}^{\text{SFT}}}$, and $\pi_{\text{ref}}$ (Li et al., 20 Aug 2025).

2. Relation to standard speculative sampling

Classical speculative sampling uses a cheap draft model $p$ and an expensive target model $q$, accepts a drafted token with probability

$$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$0

and, on rejection, samples from the residual

$$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$1

thereby preserving the target model distribution exactly, up to hardware numerics (Chen et al., 2023). Standard speculative sampling is therefore an exact acceleration method for the unshifted target distribution.

SSS preserves the draft-and-verify structure but changes the objective. If one naively replaces the standard draft with an aligned draft while keeping the usual verifier unchanged, standard speculative sampling still attempts to recover the unaligned target distribution $$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$2, not the aligned distribution $$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$3. The paper also reports a severe acceptance-rate collapse in that naive setting: for OPT-6.7B target / OPT-125M draft, standard draft acceptance is $$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$4 while aligned draft acceptance is $$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$5 ($$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$6); for OPT-13B target / OPT-350M draft, standard draft acceptance is $$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$7 while aligned draft acceptance is $$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$8 ($$\pi^\star(y|x) \propto \pi_{\text{ref}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$9) (Li et al., 20 Aug 2025).

SSS resolves this by keeping proposals from the aligned draft $\pi_{\text{draft}^r}$0 but comparing them against $\pi_{\text{draft}^r}$1 in the acceptance ratio, rather than against $\pi_{\text{draft}^r}$2 itself. This is the decisive departure from standard speculative sampling. It is not merely a reward-guided reranking layer on top of ordinary speculation; it is a shifted-target verifier construction whose correction terms are intended to recover $\pi_{\text{draft}^r}$3 rather than $\pi_{\text{draft}^r}$4 (Li et al., 20 Aug 2025).

3. Acceptance rule, bonus distribution, and decoding procedure

The SSS acceptance probability for a drafted token $\pi_{\text{draft}^r}$5 is

$\pi_{\text{draft}^r}$6

This is the core algorithmic change. Proposals are sampled from $\pi_{\text{draft}^r}$7, but acceptance is calibrated against the SFT draft and the unchanged target (Li et al., 20 Aug 2025).

When rejection occurs after $\pi_{\text{draft}^r}$8 accepted draft tokens, the replacement token is drawn from the reward-shifted bonus distribution

$\pi_{\text{draft}^r}$9

Under the aligned-draft assumption, this becomes

$\pi_{\text{ref}}$0

which clarifies its role as a residual correction relative to the RLHF-optimal target (Li et al., 20 Aug 2025).

Operationally, SSS proceeds as follows. It initializes $\pi_{\text{ref}}$1 and repeatedly samples a lookahead block of $\pi_{\text{ref}}$2 draft tokens

$\pi_{\text{ref}}$3

The target model then evaluates likelihoods for the drafted block. Each drafted token is tested with the modified acceptance criterion. Accepted tokens are appended to the output. At the first rejection, a replacement is sampled from $\pi_{\text{ref}}$4, later speculative tokens in the block are discarded, and drafting restarts from the corrected prefix (Li et al., 20 Aug 2025).

One subtle but important difference from standard speculative sampling is that SSS omits the usual “extra token” from the target model when all $\pi_{\text{ref}}$5 draft tokens are accepted. The paper states that no additional token is sampled from the target model once all draft tokens are accepted, since the objective is not the target model distribution (Li et al., 20 Aug 2025).

4. Exactness result and its assumptions

The paper’s main theorem states that, under the assumption

$\pi_{\text{ref}}$6

the probability that SSS generates response $\pi_{\text{ref}}$7 given prompt $\pi_{\text{ref}}$8 is exactly

$\pi_{\text{ref}}$9

Thus, the method is presented as an exact sampler for the RLHF-optimal reward-shifted target, not as a heuristic preference-guided decoder (Li et al., 20 Aug 2025).

The proof mirrors the standard speculative-sampling argument. A token can arise either from an accepted draft proposal or from bonus-token correction after rejection. For the accepted part,

$$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$0

The rejection mass is then

$$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$1

and the bonus distribution contributes

$$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$2

Combining these two pieces yields

$$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$3

The sequential autoregressive extension then gives the full response law (Li et al., 20 Aug 2025).

This exactness claim depends on several conditions. The paper relies on the KL-regularized RLHF form $$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$4, on the well-aligned draft assumption, on implicit shared-support conditions so that ratios such as $$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$5 are well-defined, and on a shared reward signal linking the aligned draft and the notional aligned target (Li et al., 20 Aug 2025). The paper also acknowledges that the aligned-draft assumption is difficult to verify in practice because the true reward $$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$6 is not directly known. This motivates a softened empirical variant in which the bonus term is modified to

$$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$7

with $$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$8 recovering the original SSS (Li et al., 20 Aug 2025).

The central theoretical insight is that the mismatch between the aligned draft and the unaligned target is treated as a resource rather than a liability. Because both the aligned draft and the desired aligned target contain the same exponential reward tilt, the product

$$\pi_{\text{draft}^r}(y|x) \approx \pi_{\text{draft}^{\text{SFT}}}(y|x)\cdot\exp\left(\frac{1}{\beta}r(x,y)\right).$$9

recovers the desired reward-shifted large-model distribution up to normalization. This suggests that SSS is best understood as a shifted-target speculative sampler, not as ordinary speculative decoding with a reward-biased draft.

5. Experimental evidence and efficiency profile

The reported experiments use the HH-RLHF dataset, draft–target pairs with shared vocabulary and similar distributions, and draft post-training by SFT followed by DPO. The evaluated pairs are Qwama-0.5B $\pi_{\text{draft}^r}$0 Llama-3-8B, OPT-125M $\pi_{\text{draft}^r}$1 OPT-6.7B, and OPT-350M $\pi_{\text{draft}^r}$2 OPT-13B. Gold reward is measured by a “LLaMA 7B Reward Model” trained on HH-RLHF. Inference uses temperature $\pi_{\text{draft}^r}$3, max sequence length $\pi_{\text{draft}^r}$4, and one NVIDIA L40S GPU (Li et al., 20 Aug 2025).

The main reported results are summarized below.

Target model	SSS result	Notable comparator
Llama-3-8B	Gold R $\pi_{\text{draft}^r}$5, 86.2 calls, 10.7 s, $\pi_{\text{draft}^r}$6 over BoN	BoN-10: $\pi_{\text{draft}^r}$7, 1280 calls, 58.0 s
OPT-6.7B	Gold R $\pi_{\text{draft}^r}$8, 115 calls, 7.8 s, $\pi_{\text{draft}^r}$9 over BoN	BoN-5: $\pi_{\text{draft}^{\text{SFT}}}$0, 640 calls, 22.7 s
OPT-13B	Gold R $\pi_{\text{draft}^{\text{SFT}}}$1, 112.5 calls, 13.6 s, $\pi_{\text{draft}^{\text{SFT}}}$2 over BoN	BoN-5: $\pi_{\text{draft}^{\text{SFT}}}$3, 640 calls, 69.4 s

On Llama-3-8B, SSS improves substantially over vanilla decoding and vanilla speculative decoding while using far fewer calls than BoN, TreeBoN, or CARDS; its Gold R $\pi_{\text{draft}^{\text{SFT}}}$4 remains below TreeBoN’s $\pi_{\text{draft}^{\text{SFT}}}$5 and CARDS’s $\pi_{\text{draft}^{\text{SFT}}}$6, but its latency and call profile are much smaller (Li et al., 20 Aug 2025). On OPT-6.7B, SSS attains the best gold reward among the compared methods: vanilla $\pi_{\text{draft}^{\text{SFT}}}$7, BoN-5 $\pi_{\text{draft}^{\text{SFT}}}$8, TreeBoN $\pi_{\text{draft}^{\text{SFT}}}$9, CARDS $\pi_{\text{ref}}$0, vanilla SD $\pi_{\text{ref}}$1, and SSS $\pi_{\text{ref}}$2 (Li et al., 20 Aug 2025). On OPT-13B, SSS again attains the best reported gold reward: vanilla $\pi_{\text{ref}}$3, BoN-5 $\pi_{\text{ref}}$4, TreeBoN $\pi_{\text{ref}}$5, CARDS $\pi_{\text{ref}}$6, vanilla SD $\pi_{\text{ref}}$7, and SSS $\pi_{\text{ref}}$8 (Li et al., 20 Aug 2025).

These results support two distinct claims. First, SSS can produce substantially better reward than vanilla decoding and vanilla speculative decoding. Second, it often offers a markedly better reward–latency tradeoff than BoN-style and reward-model-in-the-loop baselines, because it avoids runtime reward-model calls and still amortizes strong-model verification over speculative blocks (Li et al., 20 Aug 2025).

The paper also reports that the practical softened variant can outperform the theorem-matched $\pi_{\text{ref}}$9 setting when the aligned-draft assumption is imperfect. On OPT-13B/OPT-350M, the best performance occurs for $p$0. This indicates that theoretical exactness and empirical best performance need not coincide when the aligned draft is only approximately reward-shifted (Li et al., 20 Aug 2025).

6. Conceptual position, related methods, and limitations

SSS sits at a distinct point within the speculative-decoding literature. Standard speculative sampling is an exact accelerator for an unchanged target distribution (Chen et al., 2023). By contrast, SSS explicitly changes the target to a reward-shifted law and adjusts both acceptance and correction so that the new target, rather than the verifier distribution, is recovered (Li et al., 20 Aug 2025). This places SSS closer to a shifted-target speculative sampler than to a pure systems acceleration method.

It also differs from reward-guided speculative or speculative-like methods that do not define or recover an explicit reward-shifted target. “Fast Best-of-N Decoding via Speculative Rejection” is a sequence-level rejection and pruning scheme designed to accelerate Best-of-$p$1; it does not specify a target distribution, does not introduce a reward-shifted law such as $p$2, and has no exactness theorem of that kind (Sun et al., 2024). “Reward-Guided Speculative Decoding for Efficient LLM Reasoning” intentionally abandons exact target-distribution preservation in favor of a reward-conditioned mixture policy with process-reward gating; its guarantees are about expected reward under budget constraints, not exact recovery of a reward-tilted sequence law (Liao et al., 31 Jan 2025). “Cactus” is particularly relevant because it formalizes speculative sampling for an arbitrary target distribution $p$3 and derives acceptance/recovery rules under a divergence-constrained target shift, but its concrete shift is acceptance-optimized rather than reward-defined (Hao et al., 5 Apr 2026). This suggests that SSS belongs to the family of shifted-target speculative methods, but with a specifically RLHF-derived target.

Several common misconceptions follow from these distinctions. SSS is not simply “standard speculative sampling with an aligned draft,” because naive substitution both preserves the wrong target and sharply lowers acceptance. It is not Best-of-$p$4 acceleration, since it does not search over completed samples and rerank them with a reward model. It is also not a runtime reward-model decoding method, because the reward shift is encoded in the aligned draft and no external reward model is queried during generation (Li et al., 20 Aug 2025).

The method’s main limitations are equally explicit. It depends on a well-aligned draft satisfying

$p$5

yet the true reward function is unknown, so this assumption cannot be directly verified. Practical performance depends on draft post-training quality and on correction hyperparameters such as $p$6. Generalization is also a concern, since a small draft aligned for one task may not transfer cleanly to others (Li et al., 20 Aug 2025).

A plausible implication is that SSS is strongest when a small draft can reliably internalize the desired reward tilt and when one wants a serving-time aligner that avoids both large-model post-training and runtime reward evaluation. Under those conditions, it functions as an exact or near-exact weak-to-strong alignment mechanism at decoding time; when those conditions fail, its behavior becomes increasingly dependent on approximation quality rather than theorem-level guarantees.