AutoJudge: Judge Decoding Without Manual Annotation

TL;DR

AutoJudge addresses the latency-accuracy trade-off in large language model inference by shifting from lossless to lossy speculative decoding in a task-adaptive way. It automatically mines task-specific important token combinations through a semi-greedy search and trains a lightweight classifier to predict which mismatching tokens can be safely accepted, allowing more tokens to be drafted per verification cycle. Across mathematical reasoning (GSM8K) and programming (LiveCodeBench) tasks, and with vLLM-based inference, AutoJudge achieves up to about $2\times$ speedups with modest accuracy declines (often $\le 3\%$), without requiring manual annotation. The approach is compatible with existing speculative decoding frameworks and inference engines, enabling practical acceleration in real-world deployments, with promising extensions to EAGLE-2 and open-ended generation. Overall, AutoJudge provides a data-efficient, task-specific mechanism to accelerate inference while preserving downstream quality.

Abstract

We introduce AutoJudge, a method that accelerates large language model (LLM) inference with task-specific lossy speculative decoding. Instead of matching the original model output distribution token-by-token, we identify which of the generated tokens affect the downstream quality of the response, relaxing the distribution match guarantee so that the "unimportant" tokens can be generated faster. Our approach relies on a semi-greedy search algorithm to test which of the mismatches between target and draft models should be corrected to preserve quality and which ones may be skipped. We then train a lightweight classifier based on existing LLM embeddings to predict, at inference time, which mismatching tokens can be safely accepted without compromising the final answer quality. We evaluate the effectiveness of AutoJudge with multiple draft/target model pairs on mathematical reasoning and programming benchmarks, achieving significant speedups at the cost of a minor accuracy reduction. Notably, on GSM8k with the Llama 3.1 70B target model, our approach achieves up to $\approx2\times$ speedup over speculative decoding at the cost of $\le 1\%$ drop in accuracy. When applied to the LiveCodeBench benchmark, AutoJudge automatically detects programming-specific important tokens, accepting $\ge 25$ tokens per speculation cycle at $2\%$ drop in Pass@1. Our approach requires no human annotation and is easy to integrate with modern LLM inference frameworks.

Figures (25)

• Figure 1: Intuitive scheme of the proposed approach: (left) data collection: detecting mismatching tokens that affect final response quality; these tokens are then used to train a classifier (right) using the trained classifier to generate more tokens per cycle with speculative decoding.
• Figure 1: Inference speed benchmarks on GSM8K 0-shot with vLLM for (left) 1B draft / 8B target models with tuned window size (baseline = 8, AutoJudge = 10) and (right) for 8B draft / 70B target models (all Instruct) with tuned window size (baseline = 8, AutoJudge = 32 ).
• Figure 2: Excerpts from GSM8K (top, left) and LiveCodeBench (right) labeled by Algorithm \ref{['alg:important_tokens_mining']}. Important mismatching tokens that are in red, unimportant ones are in green. Alternative tokens are shown in [brackets]. Black tokens are where $\theta_{draft}$ and $\theta_{target}$ gave the same prediction. The top example additionally shows $\theta_{target}$ continuations after mismatching tokens ($\checkmark$ if $\alpha \equiv \hat{\alpha}$, $\bigtimes$ if not).
• Figure 2: Inference speed with vLLM for (left) Llama 3.1 8B draft / 405B target models on GSM8K 0-shot with tuned window size (baseline=14, AutoJudge=20). (right) Llama 3.1 8B draft / 70B target (all Instruct) on GSM8K 8-shot with offloading, tuned window size (baseline=10, AutoJudge=48).
• Figure 3: Accuracy and the number of accepted tokens on GSM8K for (left) 8-shot Llama-3.2 1B draft / Llama-3.1 8B target and (right) 0-shot Llama 3.1 8B draft / Llama 3.1 70B target (all Instruct)