Research map / Large language models / Reasoning and chain of thought
Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes judge decoding, which adds a linear judge head on frozen target-model embeddings so speculative decoding accepts correct but non-aligned draft tokens, yielding large speedups on Llama-3.1 models.
Problem
Standard speculative decoding verifies draft tokens by alignment with target probabilities, so many high-quality draft tokens are rejected, limiting the number of drafted tokens and speedup.
Approach
Replace or augment alignment-based verification with a learned linear judge head on the target's last-layer token embeddings, trained to predict token correctness (LLM-as-a-judge style), and OR its accept mask with standard SD's mask.
Key results
- 8B/405B-Judge reaches 9.7x speedup over standard decoding in HuggingFace and 129.3 tokens/s in gpt-fast on 8 H100s, with average acceptance length 19.7.
- Judge decoding preserves target accuracy on GSM8K, HumanEval, MT-Bench, ARC, and MMLU, while top-K verification loses most performance.
- 8B/70B-Judge reaches 141.8 tokens/s in gpt-fast, versus 76.7 for standard SD on 2 H100s.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use judge verification instead of standard speculative verification | Llama-3.1-8B draft / Llama-3.1-70B target · Llama-3.1-70B decoding | average acceptance length: 6.4 -> 18.8 | Helped |
| Use judge verification instead of standard speculative verification | Llama-3.1-8B draft / Llama-3.1-70B target · Llama-3.1-70B decoding | speedup (gpt-fast): 1.7× -> 3× | Helped |
| Use judge verification instead of standard speculative verification | Llama-3.1-8B draft / Llama-3.1-405B target · Llama-3.1-405B decoding | average acceptance length: 6.3 -> 19.7 | Helped |
| Use judge verification instead of standard speculative verification | Llama-3.1-8B draft / Llama-3.1-405B target · Llama-3.1-405B decoding | speedup (HuggingFace): 5.3× -> 9.7× | Helped |
| Use judge verification instead of standard speculative verification | Llama-3.1-8B draft / Llama-3.1-405B target · Llama-3.1-405B decoding | tokens/s: 58.7 -> 129.3 | Helped |
| Remove coding examples from judge training data (ablation) | Llama-3.1-405B judge decoding · HumanEval | accuracy: 86.6 to 80.4% | Hurt |
| Use judge decoding instead of the vanilla draft model | Llama-3.1-405B judge decoding · HumanEval | accuracy: 71.3% -> 80.4% | Helped |
| Use the simple top-K verification heuristic rather than the learned judge head | Llama-3.1-70B and Llama-3.1-405B decoding · GSM8K, HumanEval, MT-Bench, ARC, MMLU | accuracy | Hurt |
| Use embeddings from shallow layers rather than deeper layers (ablation) | Judge head · Judge correctness classification | accuracy | Hurt |
| Replace the linear judge head with MLPs or shallow Transformer networks (ablation) | Judge head · Judge correctness classification | performance | Hurt |
| Use standard alignment-based verification on high-quality draft responses | Llama-3.1-405B target with GPT-4o draft · MT-Bench, GSM8K, HumanEval | average acceptance length: roughly two | No clear effect |
Surprising findings
- Even GPT-4o as draft for Llama-405B yields only about two accepted tokens under standard verification, so alignment-based acceptance does not reward draft quality.
- Llama-405B conditioned on wrong tokens often corrects itself, so the target's token embeddings encode correctness.
Limitations
- Loses the mathematical guarantee of matching target output distribution.
- Requires a high-quality draft model; self-speculation and small drafters are not suitable.
- Out-of-distribution tasks (e.g., HumanEval when coding data excluded) require annotated in-domain data; HumanEval drops from 86.6 to 80.4.
Experimental setup
Datasets: Alpaca ARC wikipedia-summaries GSM8K HumanEval MT-Bench MMLU
Benchmarks: GSM8K HumanEval MT-Bench ARC MMLU
Metrics: average acceptance length speedup factor tokens/s accuracy
Baselines: standard speculative decoding top-K verification Medusa EAGLE-2 Llama-3.1-405B
Builds on: speculative decoding LLM-as-a-judge gpt-fast Medusa EAGLE-2
Closest papers by meaning
- AutoJudge: Judge Decoding Without Manual Annotation (NeurIPS 2025 · similarity 0.92)
- Training-Free Loosely Speculative Decoding: Accepting Semantically Correct Drafts Beyond Exact Match (ICLR 2026 · similarity 0.92)
- Online Speculative Decoding (ICML 2024 · similarity 0.92)
- SPEED-Bench: A Unified and Diverse Benchmark for Speculative Decoding (ICML 2026 · similarity 0.92)
- SelfJudge: Faster Speculative Decoding via Self-Supervised Judge Verification (ICML 2026 · similarity 0.92)
- Speculative Decoding with CTC-based Draft Model for LLM Inference Acceleration (NeurIPS 2024 · similarity 0.92)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.