Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment

Authors: Gregor Bachmann, Sotiris Anagnostidis, Albert Pumarola, Markos Georgopoulos, Artsiom Sanakoyeu, Yuming Du, Edgar Schönfeld, Ali Thabet, Jonas K Kohler
Primary area: foundation or frontier models, including LLMs
Keywords: llm inference, speculative decoding

Abstract:
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 coined TokenCourt 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 $9\times$ 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.

Decision: accept (oral)
Reviewer scores: 8, 8, 10, 6

Reviews and discussion:
<thread>
creation date:
1730315854758

thread_type:
official_review

<thread_object>
creator: reviewer_sf3q

<content>
summary:
In standard Speculative Decoding (SD), draft tokens are typically accepted or rejected during the verification phase based on whether they align with the target model distribution. However, this paper raises the concern that this verification process may be too stringent, resulting in high-quality draft tokens being rejected because they don’t align to the target distribution even though they may be valid and correct continuations. The authors try to address this concern by building upon the idea of LLM-as-a-judge and training a small judgment model that works on top of the target model latent embeddings and can be used to judge whether a draft is a correct continuation or not. Using this idea, the authors present speed ups using draft and target models in the Llama-3.1 family on a range of benchmarks and present a novel dataset (TokenCourt) used to train the target judge model.

strengths:
- The authors clearly present and motivate their ideas around standard SD verification potentially being overly stringent and distinguishing between semantic correctness and target distribution “alignedness”. The motivation is straightforward to understand with the examples provided and it’s clear how standard verification can be arbitrarily strict depending on the target model resulting in less than desired inference efficiency using standard SD at times. 
- The reported speedups over standard SD are significant, largely in part due to the increase in the number of accepted tokens using this new judge decoding verification process compared to standard verification. However, the authors correctly point out also that this is not necessarily an apples to apples comparison given the loss of output distribution guarantees compared to standard SD.
- The investigations into token embeddings signal errors was particularly interesting in that “incorrect” statements in text could be pinpointed based on the target model’s latent embedding because of the models tendency to try to immediately rectify its response following an inaccuracy and the resulting dataset curated based on this could prove useful to others if released

weaknesses:
- Given the fact that verification is being changed so there is no longer a guarantee of output distribution match, the speedup comparisons are no longer apples to apples and have to also be paired with quality comparisons to ensure quality is maintained. Additionally, given the way the judge model is trained on “correct” and “incorrect” inputs, it seems that for novel tasks, a new tuning dataset must be hand curated for each task to train the judgment model to recognize “correctness” making the overhead of applying a technique like Judge Decoding significant compared to other SD modifications.
- The notion of “correctness” is easiest to understand in the context of text generation tasks that are Question&Answer based. However, it’s unclear how this idea would work for more open-ended creative generation tasks where “correctness” may not be as easily defined and matching the target model output distribution may be ideal.
- An additional condition to applying Judge Decoding is for the draft model to be of high-quality that can produce longer valid generations otherwise the approach may not provide that much of a speedup. 
- The reported experimental results are somewhat limited such as the out of distribution performance results being restricted to only a few lines and the full results not being reported in the appendix at least in comparison to the baselines etc. Additionally, it’s unclear what task is being reported on for Table 1 whereas Figure 6 shows accuracy across a variety of tasks. It would be helpful to also see the relative speedups across each of these variety of tasks as well.
- For the most part, the authors are forthcoming about weaknesses and limitations to their judge decoding approach and many of these points will align. However, the amount of pre-requisites needed to apply the technique as well as required overhead for verification of maintained quality due to the loss of output distribution alignment are substantial.

questions:
- Do the authors plan on releasing their TokenCourt dataset?
- What happens when the target LLM isn’t one of the models used to generate the correct/incorrect responses to questions in TokenCourt used to train the linear judge module? Is this a necessary requirement for this technique to work? I ask because of a scenario where one may want to re-use a curated TokenCourt dataset for a new target model from an entirely different family of models to avoid having to hand annotate again.
- Looking at the showcased TokenCourt examples, one question is if Judge Decoding may result in potentially verbose outputs. And in those cases, even though more draft tokens are being accepted and decoded per second and the full output is semantically correct, it is not as concise or clear as the original target model output if standard verification were being applied?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730499405482

thread_type:
official_review

<thread_object>
creator: reviewer_h6mh

<content>
summary:
The authors propose "Judge Decoding," a novel approach to speculative sampling that challenges the fundamental assumption that verification must be based on model alignment. Instead of using likelihood comparisons, they train a simple linear classifier on the last embedding layer of the target model to judge token quality directly. Their key insight is that current speculative decoding methods reject many high-quality tokens simply because they don't perfectly align with the target model's preferred phrasing. By analyzing embeddings and observing how models attempt to correct errors, they show that correctness information is already present in the token representations. They leverage this insight to create a lightweight judge module (16.4k parameters) trained on their carefully curated TokenCourt dataset. The method achieves up to 9× speedups over standard decoding, reaching unprecedented speeds of 129 tokens/s for Llama-405B, while maintaining quality across a comprehensive range of benchmarks.

strengths:
- Clever idea that was well-executed. The paper is well-written and the experiments make sense. The results are impactful in practical terms (people should use this) and for spawning future research (there are many possible variations on how to train the linear classifier).

- I liked that the authors decided to go with the simple linear classifier since it worked best, rather than overcomplicating things

- I also liked the systematic investigation with illustrative examples for how judge decoding differs from speculative decoding

- Achieves state-of-the-art performance (141 tokens/s for 8B/70B-Judge, 129 tokens/s for 8B/405B) on standard hardware (2-8 H100 GPUs).

weaknesses:
I would have liked if the authors addressed the possible impact of their method on less binary (correct vs incorrect) metrics like reward model scores or chatbot ELO scores. I could imagine to get a more noticeable regression on those, since the writing style of models is more important in that setting. Perhaps something for the discussion?

The Burns et al 2022 paper (https://arxiv.org/abs/2212.03827) does some nice work on eliciting concepts with linear probes which could be relevant here.

questions:
No questions
</content>
</thread_object>

</thread>
<thread>
creation date:
1730713426422

thread_type:
official_review

<thread_object>
creator: reviewer_kfoq

<content>
summary:
The paper introduces Judge Decoding, a speculative decoding enhancement that improves token acceptance by evaluating contextual correctness rather than strict alignment with the target model. Traditional speculative decoding often rejects accurate tokens from draft models due to misalignments rather than incorrectness, limiting speedups. By training a lightweight “judge” module on a custom TokenCourt dataset, Judge Decoding accepts more valid tokens, achieving up to 9× faster generation speeds with maintained quality on the Llama-3.1 model family, demonstrating a promising approach for efficient large language model inference.

strengths:
The paper is well-written. The argument is clear and reasonable. The proposed method shows effectiveness.

weaknesses:
I do not see significant weakness.

But the sentence "To speed up inference in such a setting, Chen et al. (2023); Leviathan et al. (2023) concurrently propose speculative decoding (SD)" in the introduction section is incorrect.
At least, there is a much earlier one in 2018 that proposed speculative decoding: "Blockwise Parallel Decoding for Deep Autoregressive Models" by Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. The author skipped a large amount of literature on speculative decoding before 2023.

Authors are required to address this issue.

questions:
NA
</content>
</thread_object>

</thread>
<thread>
creation date:
1730725373004

thread_type:
official_review

<thread_object>
creator: reviewer_yijh

<content>
summary:
The submission proposes Judge Decoding (JD), a variation of Speculative Decoding (SD) for LLMs. Like SD, JD uses a small draft model to propose output, and uses a larger model to check the output. The difference between SD and JD is in how the check is performed: where SD focuses on checking that the draft model's output is exactly like the larger model's (i.e., effectively per-token equality), JD uses a weaker criterion, based on the embeddings of the final layer of the larger (judge) model. The downside of this weaker criterion is that the output of the model is not guaranteed to be equivalent to the larger model's output anymore, but experiments indicate a significant potential speedup.

strengths:
* Clear presentation, covering both technical details and high-level ideas
* Comprehensive experiments that emphasise the strengths of the proposed method
* Reported speedups seem significant, and seem like they could be very useful for practical serving of LLMs

weaknesses:
* Some of the results (e.g., lack of agreement between different models / human drafts and models) seem to be obvious and do not help to make the argument
* The use of a disjunctive criterion (standard SD accepts or JD accepts) is a bit surprising, and left me wondering in what cases (and how often) JD rejects when SD accepts. The paper would be strengthened by investigating this more closely.
* Sect. 5.3 (OOD performance) is very important for deployment of the proposed method, and so it would be useful to expand on this. In particular, Sect. 4.1/"Dataset Curation" does not go into much detail into what distribution the training data follows, and which tasks are covered well. It would be helpful to understand whether there is a strong alignment between the TokenCourt training data and the considered evals.

questions:
* Fig 2/Fig 3: It would be useful to report the number of proposed tokens as well, either explicitly or by presenting "ratio of accepted tokens" rather than an absolute number.
* One surprising result is not investigated: using Llama 405B to draft outputs of the 8B model on GSM8K (lower right of Fig. 2) seems to work less well than the other way around, even though the models are obviously well-aligned. I would have expected a basically equal result here (following the intuition that the larger model is effectively a superset of the smaller model), and was very surprised to see a lower score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732201254027

thread_type:
public_comment

<thread_object>
creator: ~weilin_zhao1

<content>
title:
Doubts About Baseline Speed in Table 1 Under the GPT-Fast Framework

comment:
Ideally, framework changes should affect all methods **consistently**. This is because changing the framework is indeed altering the runtime difference between large and small models. **However**, the impact of switching frameworks on the speedup ratio is **not consistent** across the table.

Comparing Huggingface with GPT-Fast in Table 1

1) For the 405B model, STANDARD (5.3x `>` 1.7x), JUDGE (9.7x `>` 3.9x) and Medusa (6x `>` 1.9x) show similar trend, which shows consistency.
2) For the 70B model , STANDARD (1.5x `<` 1.7x), JUDGE (2.0x `<` 3.0x) follows the same trend but EAGLE (3.3x `>` 1.9x) does not, which is inconsistent.

The authors explain in the paper that EAGLE's slowdown in GPT-Fast is similar to the findings presented in EAGLE's original paper. However, we consider it is due to the suboptimal implementation of EAGLE in GPT-Fast.

To clarify, we suggest adding these baselines for the 70B model:

1. MEDUSA: In Table 1, Medusa is only reported in the 405B model but not reported in the 70B model.

2. STANDARD speculative decoding + lenience factor: [1] has already proposed using a lenience factor to relax the acceptance criteria for speculative sampling, which aligns with the motivation of this paper and should be included as a baseline.

[1] Leviathan, Yaniv, Matan Kalman, and Yossi Matias. "Fast inference from transformers via speculative decoding." International Conference on Machine Learning. PMLR, 2023.
</content>
</thread_object>

</thread>
<thread>
creation date:
1744678907522

thread_type:
public_comment

<thread_object>
creator: ~drew_jin1

<content>
title:
Inquiry About Open-Sourcing the Code

comment:
Dear Authors,

I hope this message finds you well. I recently had the opportunity to go through your work on speculative decoding, and I must say it is truly impressive. 

I was wondering if there are any plans to **open-source the code** related to this project. It would be incredibly beneficial for the community to explore, learn from, and potentially build upon your work. If the code is already available, could you kindly point me in the right direction?

Thank you for your time, and I look forward to hearing from you.

Best regards
</content>
</thread_object>

</thread>


Metareview:
metareview:
Summary: Different from the checking mode in the standard speculative sampling framework, this paper studies the problem of accelerating LLM inference by LLM-as-a-judge: asking the LLM itself whether the drafted content is valid or not. Experiments on Llama verify the effectiveness of the proposed method.

Strength:
1. The paper is well-written.
2. The experiments are sufficient.
3. The proposed method is effective by the experiments.

Weakness:
1. For tasks beyond QA, there might be no standard answer for the correctness. It is unclear how to use LLM-as-a-judge for such tasks, as there might be multiple "correct" answers.
2. The proposed framework changes the output text distribution of the target model. So it is not directly comparable with other distribution-preserving method, such as EAGLE and Medusa. The claimed speedup ratio might not be the only criteria to evaluate the effectiveness. One needs to take into account the trade-off between quality loss and speedup ratio.
3. LLM may have hallucination. LLM-as-a-judge might not always be reliable.

Reviewers are all positive about the paper. Therefore, AC are happy to recommend the paper for acceptance with spotlight. However, AC would not recommend the paper with an oral presentation given the above limitations, mostly raised by Reviewer Sf3Q.