Kati8/ICLR_2023_2025 huggingface dataset:
authors: Gregor Bachmann
authors: Sotiris Anagnostidis
authors: Albert Pumarola
authors: Markos Georgopoulos
authors: Artsiom Sanakoyeu
authors: Yuming Du
authors: Edgar Schönfeld
authors: Ali Thabet
authors: Jonas K Kohler
reviews: 
Comments: 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.
Recommendation: 10
Reviewer confidence: 4
Soundness/correctness: 3
Clarity: 4
reviews: 
Comments: 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.
Recommendation: 8
Reviewer confidence: 4
Soundness/correctness: 3
Clarity: 4
reviews: 
Comments: 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.
Recommendation: 8
Reviewer confidence: 3
Soundness/correctness: 4
Clarity: 4
reviews: 
Comments: This work introduces 'Judge Decoding,' an innovative technique for speculative sampling that fundamentally rethinks the role of verification by moving away from the conventional reliance on model alignment metrics. Rather than employing likelihood-based comparisons, the authors develop a lightweight linear classifier that operates on embeddings from the target model's final layer to directly assess token quality. The central contribution is demonstrating that many tokens rejected by existing speculative decoding methods are actually of high quality—they simply deviate from the target model's preferred output distribution. Through careful analysis of embedding spaces and examination of model error-correction behaviors, they provide evidence that semantic correctness is encoded within token representations themselves. They operationalize this finding by constructing a compact judge module (16.4k parameters) trained on their TokenCourt dataset. The approach delivers remarkable speedups of up to 9× compared to baseline decoding, achieving exceptional throughput of 129 tokens/s for Llama-405B while preserving output quality across diverse evaluation benchmarks.
Strengths: - The core insight is innovative and the execution is rigorous. The manuscript is clearly written with well-designed experiments that validate the approach. The results have substantial practical implications (practitioners should adopt this method) and open important research directions (numerous possibilities exist for refining classifier training strategies).

- The authors demonstrated good judgment by selecting the simple linear classifier when it outperformed more complex alternatives, avoiding unnecessary complexity.

- The systematic analysis with concrete examples effectively illustrates how judge decoding diverges from traditional speculative decoding methodologies.

- Demonstrates cutting-edge throughput results (141 tokens/s for 8B/70B-Judge, 129 tokens/s for 8B/405B) achievable with commodity compute resources (2-8 H100 GPUs).
Weaknesses: It would strengthen the paper if the authors examined potential performance degradation on continuous evaluation metrics such as reward model scores or ELO ratings from pairwise comparisons. The method may show more pronounced quality drops on these metrics since stylistic alignment with the target model becomes more significant in these evaluation frameworks. This discussion would be a valuable addition to the manuscript.

The work by Burns et al. (2022) (https://arxiv.org/abs/2212.03827) on extracting interpretable features via linear probes appears methodologically relevant and warrants consideration.
Recommendation: 10
Reviewer confidence: 4
Soundness/correctness: 3
Clarity: 4