Introducing AutoJudge: Streamlined inference acceleration via automated dataset curation

Speculative decoding speeds up generation by pairing a small “draft” model with the main “target” model. The draft proposes several next tokens; the target runs in parallel to verify them. Tokens that match the target are accepted; the first mismatch (and everything after) is rejected. This keeps the output distribution identical to the target model’s own decoding.

In practice, strict distribution matching isn’t always necessary. Lossy variants trade a tiny amount of quality for more speed. Judge Decoding is one such approach: it only rejects a mismatch if accepting it would harm the final answer. For example, math errors or logical bugs matter, but minor stylistic changes often don’t. Our work builds directly on this idea.

The catch with Judge Decoding is data: it needs humans to label “critical” mismatching tokens for each task, which is expensive and doesn’t transfer perfectly across domains. AutoJudge removes this bottleneck by automatically mining those important tokens—no human annotators required. 

The AutoJudge Method

AutoJudge consists of the following stages:

1. Automatically mine “important” mismatches
   For a prompt, generate a target answer and locate where draft and target models disagree. Iteratively swap draft ↔ target tokens and re-evaluate the task (e.g. GSM8K answer equality or code unit tests). A mismatch is important if keeping the draft token breaks the final answer; otherwise it’s unimportant. This semi-greedy pass reliably surfaces at least one important token whenever answers differ.
2. Train a tiny classifier on existing embeddings
   We use a simple logistic regression fed by transformer hidden states already computed during speculative decoding. Concatenating draft and target token embeddings works best and remains robust across regularization choices and small architectural variants.‍
3. Accept “unimportant” mismatches at verification time
   During the verification phase—exactly where the baseline would reject a mismatching draft token—we call the classifier. If it predicts that the token is unimportant, we accept it and keep moving forward, increasing accepted tokens per speculation cycle. The approach is compatible with standard, tree-based, and single-model multi-head speculative decoding methods, and slots into popular stacks like vLLM, TensorRT-LLM, and TGI. In practice, we pick a high-recall threshold (≥90%) to safeguard accuracy while still skipping a large fraction of tokens.

In Figure 2, we demonstrate how AutoJudge can accept more tokens during the speculative decoding step. AutoJudge adds a tiny “judge” that asks, at each mismatch, whether the difference actually changes the final answer. In the example, the mismatch is a harmless wording — like “equals” vs “becomes” — and we accept it and keep the rest of the drafted tokens. If it would change correctness — like “+” vs “−” in a math step — we reject it. By only rejecting critical mismatches, we keep longer chunks from the draft, so more tokens are accepted at once and generation is faster with little impact on quality.

Performance benchmarks

Accuracy-acceptance tradeoffs

In Figure 3 we show how AutoJudge shifts the speed–quality frontier: as the number of accepted draft tokens per cycle increases (x-axis), AutoJudge (red) stays near the accuracy of lossless speculative decoding while accepting more draft tokens, unlike a naive Top-K baseline whose accuracy drops quickly. This holds in both model pairs (1B/8B on the left, 8B/70B on the right), so you can choose a threshold that yields higher tokens/s with minimal accuracy cost. In Figure 3 (right), we show that it is possible to safely accept three times more tokens, paying only a 1% accuracy drop, demonstrating that speculative decoding can safely accept up to 45 tokens with minimal loss in quality.

Inference speedup

We integrated AutoJudge into vLLM’s speculative decoding and measured end-to-end tokens/s on GPUs. (Setups included A100/H100; see notes below for model pairs.)

Mathematical reasoning (GSM8K)

Across model pairs, AutoJudge delivers consistent throughput gains with small accuracy trade-offs:

• Llama-3.1-405B (target) / 8B (draft), 8xH100: 91.5% (≈4% drop), 60.1 tokens/s, 1.20×.
• Llama-3.1-70B (target) / 8B (draft), 4xA100: 89.9% (≈2% drop), 107.4 tokens/s, 1.49×.
• Llama-3.1-8B (target) / 1B (draft), 1xA100: 80.2% (≈3% drop), 169.2 tokens/s, 1.14×.
  Baselines: 50.0 (405B), 72.3 (70B), 147.7 (8B) tokens/s.

Programming (LiveCodeBench)

AutoJudge automatically identifies critical tokens for code and boosts acceptance rates:

• Llama-3.1-70B (target) / 8B (draft): Pass@1 28.0% (≈3% drop), ~35 accepted tokens/cycle (≈3.5×). Baseline accepted tokens: ~10
  
• Llama-3.1-8B (target) / 1B (draft): Pass@1 14.5% (≈2.5% drop), ~30 accepted tokens/cycle (≈2.3×). Baseline accepted tokens: ~13

Offloading scenarios (bandwidth-limited)

When the link bandwidth is the bottleneck, longer draft windows become viable and speedups amplify:

• 8B → 70B (GSM8K): 2.4 tokens/s, 1.98×, accuracy 90.4% (≈3% drop).
• 8B → 70B (GSM8K): 1.4 tokens/s, 1.20×, accuracy 95.4% (≈+0.5%).
  Baseline: 1.19 tokens/s.
  
  

Composing with EAGLE-2

AutoJudge stacks with EAGLE-2 (which drafts from the target’s hidden states, no separate draft model). On GSM8K (0-shot) with Llama-3.1-8B-Instruct, AutoJudge adds ~8–20% tokens/s over EAGLE at small accuracy deltas: 96.8, 102.6, 107.5 tokens/s vs 89.8 baseline, with accuracies 81.3%, 81.0%, 78.1%.

Limitations & practical notes

• Speedups depend on how often mismatches are genuinely unimportant for the metric (e.g., answer equality, unit tests). Tasks such as creative writing often leave less headroom; see more experiments (including long-context GSM and writing) in the paper appendix.
• One can favor high-recall classifier thresholds (≥90%) to protect quality while still skipping many tokens. Threshold values should ideally be tuned per task.

Conclusion

AutoJudge offers a simple and fully automated algorithm that accelerates the speculative decoding loop: accept harmless mismatches, save target model calls, and go faster. It removes manual labeling from judge-style methods, learns what matters per task, and uses a tiny classifier on embeddings you already compute to ensure low runtime overhead. 

Try it

• Paper: arxiv.org/abs/2504.20039
• Code: github.com/garipovroma/autojudge 
• Data: https://huggingface.co/datasets/mightyneighbor/AutoJudge 

References‍

[1] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding, 2023. URL https://arxiv.org/abs/2211.17192

[2] Gregor Bachmann, Sotiris Anagnostidis, Albert Pumarola, Markos Georgopoulos, Artsiom Sanakoyeu, Yuming Du, Edgar Schönfeld, Ali Thabet, Jonas Kohler. Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment, 2025. URL: https://arxiv.org/abs/2501.19309 

[3] Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle-2: Faster inference of language models with dynamic draft trees, 2024. URL https://arxiv.org/abs/2406.16858