Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Revisiting Judge Decoding from First Principles via Training-Free Distributional Divergence

Shengyin Sun, Yiming Li, Renxi Liu, Weizhe Lin, Hui-Ling Zhen, Xianzhi Yu, Mingxuan Yuan, Chen Ma

TL;DR

It is revealed that the ``criticality''scores learned via costly supervision are intrinsically encoded in the draft-target distributional divergence, and a simple, training-free verification mechanism based on KL divergence is proposed.

Abstract

Judge Decoding accelerates LLM inference by relaxing the strict verification of Speculative Decoding, yet it typically relies on expensive and noisy supervision. In this work, we revisit this paradigm from first principles, revealing that the ``criticality'' scores learned via costly supervision are intrinsically encoded in the draft-target distributional divergence. We theoretically prove a structural correspondence between learned linear judges and Kullback-Leibler (KL) divergence, demonstrating they rely on the same underlying logit primitives. Guided by this, we propose a simple, training-free verification mechanism based on KL divergence. Extensive experiments across reasoning and coding benchmarks show that our method matches or outperforms complex trained judges (e.g., AutoJudge), offering superior robustness to domain shifts and eliminating the supervision bottleneck entirely.

Revisiting Judge Decoding from First Principles via Training-Free Distributional Divergence

TL;DR

It is revealed that the ``criticality''scores learned via costly supervision are intrinsically encoded in the draft-target distributional divergence, and a simple, training-free verification mechanism based on KL divergence is proposed.

Abstract

Judge Decoding accelerates LLM inference by relaxing the strict verification of Speculative Decoding, yet it typically relies on expensive and noisy supervision. In this work, we revisit this paradigm from first principles, revealing that the ``criticality'' scores learned via costly supervision are intrinsically encoded in the draft-target distributional divergence. We theoretically prove a structural correspondence between learned linear judges and Kullback-Leibler (KL) divergence, demonstrating they rely on the same underlying logit primitives. Guided by this, we propose a simple, training-free verification mechanism based on KL divergence. Extensive experiments across reasoning and coding benchmarks show that our method matches or outperforms complex trained judges (e.g., AutoJudge), offering superior robustness to domain shifts and eliminating the supervision bottleneck entirely.
Paper Structure (32 sections, 10 theorems, 10 equations, 15 figures, 1 table)

This paper contains 32 sections, 10 theorems, 10 equations, 15 figures, 1 table.

Table of Contents

  1. Introduction
  2. Background
  3. Speculative Decoding
  4. Judge Decoding
  5. The Bottlenecks of Supervision in Judge Decoding
  6. Judge Decoding With Manual Annotation
  7. Judge Decoding Without Manual Annotation
  8. Can We Bypass Data Annotation?
  9. Empirical Perspective
  10. Theoretical Perspective
  11. Experiment
  12. Experimental Setup and Baselines
  13. Benchmarks and Models.
  14. Baselines and Metrics.
  15. Main Results
  16. ...and 17 more sections

Key Result

Theorem 4.1

The empirical alignment between KL-based thresholding and the trained linear classification stems from their shared dependence on the primitives $\Delta_{ij}(x)$: (i) KL Divergence as Quadratic Aggregation: Under a second-order expansion, the KL divergence acts as a weighted quadratic sum of the pri

Figures (15)

  • Figure 1: Relationship between AutoJudge score and token-level KL divergence. Higher AutoJudge scores indicate more critical tokens, which coincide with larger KL divergence and stronger target–draft disagreement (Case 1: inconsistent key numbers). Conversely, low scores correspond to small KL divergence and minor, non-semantic deviations (Case 2: capitalization only).
  • Figure 2: Judge decoding as learning distributional discrepancy. (a) Manual-annotation pipeline. (b) Divergence-point mining without manual labels; (c) We show that a linear judge’s score is empirically correlated with and theoretically connected to distributional divergence (e.g., KL divergence).
  • Figure 3: Illustration of overextended labeling boundaries in manual annotation, adapted from DBLP:conf/iclr/BachmannAnagnostidis25. Key: Non-critical Tokens, Critical Tokens, Logic-pivoting Tokens. The span-based labeling (red) obscures the true logic-pivoting tokens (red underline), creating noisy supervision signals.
  • Figure 4: Efficiency analysis of AutoJudge dataset mining. The substantial growth in GPU hours and memory footprint for larger models highlights a significant scalability bottleneck.
  • Figure 5: Consistency analysis of critical tokens. The low percentage of consistently critical tokens (Level 4) indicates that heuristic mining is heavily influenced by generation randomness.
  • ...and 10 more figures

Theorems & Definitions (19)

  • Theorem 4.1: Structural Correspondence of KL and Linear Classifiers
  • proof
  • proof : Proof Sketch of Theorem \ref{['ssy1122:structural_correspondence']}
  • Lemma A.1: KL as a Bregman divergence
  • proof
  • Lemma A.2: Second-order expansion around $z_d$
  • proof
  • Corollary A.3: KL's quadratic form controlled by a linear map of the concatenated state
  • Lemma A.4: Pairwise decomposition of the Fisher quadratic form
  • proof
  • ...and 9 more