Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment

Gregor Bachmann; Sotiris Anagnostidis; Albert Pumarola; Markos Georgopoulos; Artsiom Sanakoyeu; Yuming Du; Edgar Schönfeld; Ali Thabet; Jonas Kohler

Source-linked AI summary

Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment

Gregor Bachmann, Sotiris Anagnostidis, Albert Pumarola, Markos Georgopoulos, Artsiom Sanakoyeu, Yuming Du, Edgar Schönfeld, Ali Thabet, Jonas Kohler

arXiv:2501.19309v1Jan 31, 2025cs.LGcs.CL

TL;DR

Large-model inference is costly, while standard speculative decoding rejects many objectively correct draft tokens because verification prioritizes alignment with the target response. The paper adds a judge-inspired embedding module to recognize token quality, achieving substantial speedups while preserving target-model quality across benchmarks.

Problem
Standard speculative decoding frequently rejects correct continuations because its alignment-based verification does not reflect contextual quality, limiting acceptance and speedup.
Method
The paper trains a compact judge module on target token embeddings and combines its acceptance decisions with standard speculative-decoding verification.
Results
Up to 9× speedup and 129 tokens/s are achieved for Llama-405B while nearly preserving target performance across benchmarks.
Takeaways & Limitations
Judge decoding enables modern small language models to draft more tokens without loss of target quality, with benefits extending to optimized inference frameworks.
Takeaways & Limitations
Standard speculative decoding remains inefficient because large portions of correct answers can be rejected when verification relies on target alignment.

Abstract

from arXiv · show

The performance of large language models (LLMs) is closely linked to their underlying size, leading to ever-growing networks and hence slower inference. Speculative decoding has been proposed as a technique to accelerate autoregressive generation, leveraging a fast draft model to propose candidate tokens, which are then verified in parallel based on their likelihood under the target model. While this approach guarantees to reproduce the target output, it incurs a substantial penalty: many high-quality draft tokens are rejected, even when they represent objectively valid continuations. Indeed, we show that even powerful draft models such as GPT-4o, as well as human text cannot achieve high acceptance rates under the standard verification scheme. This severely limits the speedup potential of current speculative decoding methods, as an early rejection becomes overwhelmingly likely when solely relying on alignment of draft and target. We thus ask the following question: Can we adapt verification to recognize correct, but non-aligned replies? To this end, we draw inspiration from the LLM-as-a-judge framework, which demonstrated that LLMs are able to rate answers in a versatile way. We carefully design a dataset to elicit the same capability in the target model by training a compact module on top of the embeddings to produce ``judgements" of the current continuation. We showcase our strategy on the Llama-3.1 family, where our 8b/405B-Judge achieves a speedup of 9x over Llama-405B, while maintaining its quality on a large range of benchmarks. These benefits remain present even in optimized inference frameworks, where our method reaches up to 141 tokens/s for 8B/70B-Judge and 129 tokens/s for 8B/405B on 2 and 8 H100s respectively.

1 INTRODUCTION

Large models improve performance but make autoregressive inference increasingly costly. Standard speculative decoding accelerates generation by verifying draft tokens against the target, yet alignment-based rejection limits acceptance of high-quality continuations; the paper proposes judge-inspired verification to address this gap.

Motivation: Larger language models often achieve better benchmark results, but their size increases deployment resource requirements and slows inference.Autoregressive decoding is especially memory-bound at small batch sizes.
Motivation: Speculative decoding uses a fast draft model to propose M candidate tokens that the target model verifies in parallel.Parallel processing can add no latency in the memory-bound regime.
Problem: Standard verification judges tokens by target-model alignment rather than contextual quality, making early rejection overwhelmingly likely as M grows.Current approaches therefore use small candidate lengths, typically M ∈{5, 7}.
Approach: LLM-as-a-judge motivates adapting verification to recognize correct but potentially unaligned continuations.The paper trains a compact module on target embeddings to produce judgements of the current continuation.
Contributions: 9× speedup over standard decoding and 129 tokens/s are reported for Llama 8B/70B-Judge while maintaining Llama-405B quality across benchmarks.The contribution also reports a linear layer trainable in under 1.5 hours.

2 RELATED WORKS

Prior speculative-decoding research has explored specialized or smaller draft models and shallow embedding-based predictors. Other work improves acceptance through distillation, token trees, or greater information exchange between draft and target models.

Draft models: Early speculative-decoding methods use specialized draft models or smaller target-model versions trained with matching data and protocols.A separate line uses shallow networks over target embeddings to predict multiple future tokens.
Acceptance improvements: Several methods improve acceptance by encouraging alignment, constructing top-K token trees for parallel verification, or exchanging more information between draft and target.Tree-based approaches use tree attention to cover more token combinations.

3 VERIFICATION IN SPECULATIVE DECODING

Speculative decoding drafts candidate tokens with a smaller model, verifies them in parallel using the target model, and accepts them sequentially under a probability-based rule. Standard verification often rejects contextually correct tokens, so accepted-token counts saturate as drafting length increases and high-quality drafts still face substantial rejection.

Background: Speculative decoding uses a draft model to autoregressively produce M candidate tokens, which the target model processes in parallel.The target produces probability vectors for the candidates and an additional next-token vector.
Background: A candidate is accepted only if all preceding candidates were accepted and a stochastic probability test succeeds.At least one valid target token is always produced, including after the first rejection or when all candidates are accepted.
Number of draft tokens: The number of accepted tokens quickly saturates as draft length M increases for Llama-3.1-8B/405B on MT-Bench and GSM8K.Consequently, larger M lowers acceptance rates and spends more drafting effort on rejected candidates.
Rejected tokens: Even competitive Llama-8B drafts frequently lose tokens despite producing accurate solutions, because standard verification seeks alignment with the target response rather than contextual correctness.The target can reject a correct answer after only two tokens.
High-quality drafts: GPT-4o drafts and high-quality human-written text also receive substantial rejection under standard verification with Llama-405B.These experiments test whether higher candidate quality improves acceptance rates.
Conclusion: Standard verification is therefore highly inefficient because large portions of correct answers are rejected, motivating verification based on quality rather than model alignment.The paper introduces a more effective scheme in the following section.

4 JUDGE DECODING

Judge decoding adapts speculative verification to assess token quality beyond target-model alignment. It uses target embeddings, a trained judge head, and an OR rule with standard verification to accept more valid draft tokens.

LLM-as-a-judge motivates evaluating candidate responses for correctness and contextual quality rather than exact alignment with the target.
Target-model embeddings of erroneous tokens flag errors and contradictions, suggesting they can support token-level correctness judgments.The target may generate corrective continuations after processing an incorrect token.
The training set contains 500 manually reviewed question, correct-answer, wrong-answer tuples with token labels marking errors as negative examples.Correct-answer tokens and wrong-answer tokens before the mistake are positive; mistaken tokens are negative, yielding roughly 20× more positive than negative examples.
A weighted-cross-entropy-trained linear judge head is placed on target embeddings, emphasizing negative examples and using early stopping to reduce false acceptance and overfitting.
Judge decoding accepts a candidate when either standard verification or the judge head accepts it, using a sigmoid confidence threshold δ that defaults to 0.5.Setting δ = 1 reduces the method to standard speculative decoding.
Figure 5 depicts forced incorrect response tokens in red, freely generated corrective tokens in green, and candidate-token judgements produced from embeddings.

5 EVALUATION OF JUDGE DECODING

Judge decoding preserves target-model accuracy while accepting substantially longer continuations and delivering strong speedups, including in optimized inference settings. Its benefits are strongest when target verification is expensive, though performance depends on task-matched judge training.

5 EVALUATION OF JUDGE DECODING: GPT-4o and human-generated continuations receive more accepted tokens under judge decoding than under standard speculative decoding across datasets.An example shows judge decoding fully accepting a correct GPT-4o response that standard speculative decoding rejects after two words.
5.1 PERFORMANCE BENCHMARK: Judge decoding almost exactly preserves target performance across benchmarks while accepting up to ∼20 draft tokens on average.The learned judge substantially outperforms top-K verification, which remains close to draft-model quality even with K = 5.
5.2 SPEED BENCHMARK: Optimized inference frameworks amplify judge decoding’s advantage by increasing the latency gap between the draft and target models.For Llama-70B, the method outperforms current state-of-the-art methods in the optimized setting, despite a smaller advantage in simple frameworks.
5.2 SPEED BENCHMARK: 9.7× speedup and 129 tokens/s are achieved for 8B/405B-Judge in HuggingFace and gpt-fast, respectively.The method accepts close to 20 tokens at a time, whereas Medusa and Eagle-23 are limited to drafting six or fewer.
5.3 OUT-OF-DISTRIBUTION PERFORMANCE: On HumanEval, out-of-distribution evaluation drops from 86.6% to 80.4% but remains above the draft model’s 71.3%.The result indicates partial transfer of correctness across tasks, while preserving target quality requires training on data of similar nature.

6 CONCLUSION

The paper identifies alignment-based rejection as a central limitation of speculative decoding and proposes judge-based verification to accept objectively correct continuations. This enables speedups up to 9× and 129 tokens/s while maintaining quality on standard benchmarks, but removes the exact target-quality guarantee.

6 CONCLUSION: Alignment-focused verification rejects objectively correct continuations, limiting speculative decoding’s ability to exploit improving small draft models.The paper frames this rejection behavior as the key limitation motivating judge decoding.
6 CONCLUSION: Judge-based verification enables drafting more tokens and achieves speedups up to 9×, including 129 tokens/s for Llama-405B.The paper reports maintained quality on standard benchmarks with a well-trained judge.
6 CONCLUSION: Judge decoding loses the mathematical guarantee of matching target quality, although extensive experiments show no quality loss on standard benchmarks.The conclusion presents the approach as a significant first step despite this trade-off.

A LIMITATIONS

The approach has explicit scope boundaries: it cannot guarantee target quality on novel tasks, depends on capable draft and target models, and requires task-relevant judge data. Safety risks from unsafe draft tokens are also not thoroughly investigated.

A LIMITATIONS: Judge decoding loses the mathematical guarantee of matching target quality, with no certainty on novel tasks.The authors report maintained quality experimentally but distinguish this evidence from a formal guarantee.
A LIMITATIONS: The method depends on a high-quality draft model and a sufficiently large target model for useful acceptance decisions.Small drafters deteriorate when projecting farther ahead, while speedups are harder to achieve for smaller target models such as Llama-8B.
A LIMITATIONS: New tasks require careful annotation of task-relevant data to maintain quality, although the required amount is small in this case.The limitation concerns transfer to tasks unlike those represented in judge training.
A LIMITATIONS: Unsafe draft tokens could be accepted by the judge, creating potential safety-critical outputs that the target would otherwise never produce.The authors did not observe this issue but also did not investigate it thoroughly because it was beyond the work’s scope.

B.1 LINEAR HEAD

The verification head is a linear module whose dimensions scale with the target model, and it adds practically negligible inference overhead by functioning as an additional vocabulary entry.

B.1 LINEAR HEAD: The linear head has dimension 16,384 for Llama-405B and 8,192 for Llama-70B.It is trained with AdamW using learning rate η = 0.0001, weight decay 0.1, and batch size 128.
B.1 LINEAR HEAD: The linear head adds practically zero inference overhead because it can be viewed as an additional entry in the vocabulary V.

B.2 JUDGE MASKING

Judge decoding combines a judge-derived decision mask with standard speculative-decoding verification. When the two masks disagree, later tokens are rejected automatically, preserving the standard correction behavior while reducing unnecessary steps.

Judge decoding combines its decision mask with the standard speculative-decoding mask.The combined masking procedure is illustrated in Figure 7.
When judge and standard verification disagree, subsequent tokens are rejected automatically, even if individually accepted.
In rare cases where judge decoding rejects a token that the target would reproduce, combining both masks still produces the same reply in fewer steps.This situation can recur across consecutive tokens.
Experiments run on H100-SXM5 GPUs, using eight GPUs and 8-bit quantization for Llama-405B and two GPUs for Llama-70B.

C MORE PROMPTS

Additional examples show judge decoding accepting long correct continuations that standard verification rejects early, while also identifying incorrect draft responses and subtle reasoning or coding errors.

C.1 REJECTED REPLIES FOR LLAM A-8B: C.1 shows Llama-8B giving completely correct answers that Llama-405B rejects early under standard verification.The examples also include wrong Llama-8B responses that judge decoding rejects rather than blindly accepting.
C.1 REJECTED REPLIES FOR LLAM A-8B: Judge decoding accepts a long continuation for a correct response that standard decoding rejects early, although it later rejects over-cautiously.
C.1 REJECTED REPLIES FOR LLAM A-8B: For an incorrect roommate-count response, judge decoding rejects the answer later and correctly identifies the error.The example concerns a situation with five roommates in total.
C.2 JUDGING OF WIKIPEDA ARTICLES: C.2 compares greedy matching in standard speculative decoding with judge-based verification of Wikipedia-summary responses.Prompts cover topics including aluminium, Moore’s Law, and Pet Shop Boys.
C.3 FORCING WRONG REPLIES FOR LL A M A-405B: Llama-405B can also acknowledge that its completion is wrong when it cannot repair the response.The paper presents this behavior as evidence that correctness may be detectable in token embeddings.
C.3 FORCING WRONG REPLIES FOR LL A M A-405B: C.3 shows Llama-405B correcting a mistaken arithmetic response by subtracting l[i] to double rather than triple.
C.3 FORCING WRONG REPLIES FOR LL A M A-405B: A coding example shows Llama-405B detecting an undefined index variable j and a subtler list-mutation bug caused by omitting a copy.