Chi Xie
Author directory2026
Calibrated Tree-Structured Speculative Decoding: Mitigating Overconfidence via Path-Level Calibration
Zixian LI | Xiaohui Song | Chi Xie | Ke Chen | Chen Chen | Haonan Lu
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Zixian LI | Xiaohui Song | Chi Xie | Ke Chen | Chen Chen | Haonan Lu
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Tree-structured speculative decoding accelerates large language model inference by drafting multiple candidate tokens in parallel and verifying them against a target model. Dynamic tree expansion strategies further improve efficiency by using draft model confidence to prune unlikely branches and stop expansion early. However, we find that existing draft models exhibit severe overconfidence: they assign high probabilities to tokens that will be rejected, causing unreliable early stopping and wasted computation on incorrect branches.
We propose Calibrated Tree-Structured Speculative Decoding (CATS-SD), a training framework for tree-based confidence calibration. Our approach calibrates draft model confidence at two complementary levels: (1) token-level calibration that aligns draft-target token ranking, and (2) path-level calibration that supervises confidence over tree-structured decoding paths. We also introduce a calibration metric using Pearson correlation to quantify confidence-acceptance alignment. On Qwen2.5-3B-Instruct, CATS-SD achieves up to 11.3% higher acceptance length and 4.0% better confidence-acceptance correlation compared to the baseline, with end-to-end speedup improving from 2.51× to 2.71× on GSM8K. We validate these results across mathematical reasoning (GSM8K), multi-turn conversation (MT-bench), and code generation (HumanEval).
We propose Calibrated Tree-Structured Speculative Decoding (CATS-SD), a training framework for tree-based confidence calibration. Our approach calibrates draft model confidence at two complementary levels: (1) token-level calibration that aligns draft-target token ranking, and (2) path-level calibration that supervises confidence over tree-structured decoding paths. We also introduce a calibration metric using Pearson correlation to quantify confidence-acceptance alignment. On Qwen2.5-3B-Instruct, CATS-SD achieves up to 11.3% higher acceptance length and 4.0% better confidence-acceptance correlation compared to the baseline, with end-to-end speedup improving from 2.51× to 2.71× on GSM8K. We validate these results across mathematical reasoning (GSM8K), multi-turn conversation (MT-bench), and code generation (HumanEval).