UBTree: Parallel Tree Drafting via Unigram and Bigram Models for Speculative Decoding

Chumeng Liang; Linxuan Wang; Xinyu Peng; Huabin Liu; Yuxin Chen; Ge Liu; Guang Lin; Qifan Song

cs.CLSep 30, 2026

UBTree: Parallel Tree Drafting via Unigram and Bigram Models for Speculative Decoding

Authors: Chumeng Liang, Linxuan Wang, Xinyu Peng, Huabin Liu, Yuxin Chen, Ge Liu, Guang Lin, Qifan Song,

Organizations: Inclusion AI · University of Illinois Urbana-Champaign · Purdue University

Abstract

Speculative decoding accelerates language model inference by verifying multiple draft tokens in a single target-model pass. Recent parallel drafters have achieved breakthrough performance in frontier production models, but their effectiveness deteriorates as the entropy of target distributions increases due to insufficient draft diversity. To overcome this bottleneck without sacrificing parallelism, we introduce UBTree, a parallel drafter that couples a Unigram proposer with a Bigram selector to construct drafting Trees. The unigram proposer is trained with the standard cross-entropy objective to generate candidate tokens independently for each position, while a lightweight bigram selector predicts transition scores between adjacent candidate pairs. Unlike the proposer, the selector is trained with a renormalized KL objective on high-temperature data. This tree-native training broadens the supervision beyond the greedy path, encouraging plausible alternative branches that improve the chance of accepting additional tokens during tree verification. Across seven standardized benchmarks with Qwen3-4B and Qwen3-8B, UBTree achieves an average speedup of $5.84$ -- $6.94\times$ over autoregressive decoding and outperforms DARTree in all 28 comparisons. Production-scale evaluation further demonstrates UBTree's advantage over frontier baselines such as DSpark.

Figures & tables

Figure 1: (a) Speedups on Qwen3-8B: Speedup over autoregressive decoding across seven benchmarks on Qwen3-8B instruct at

T_{\mathrm{infer}}=1

. (b) Speedups on Qwen3.8-27B: Serving speedup in SGLang on Qwen3.8-27B at

T_{\mathrm{infer}}=1

. (c) Long-context acceptance: Acceptance length

\tau

on Ling3-Flash-124B for LongBench-v2 and SWE-bench over different context lengths. UBTree outperforms competing methods across most benchmarks and model scales.

Figure 2: Overview of UBTree. Left: UBTree Architecture. The unigram proposer generates top- $K$ candidates at each draft position in parallel. The bigram selector scores adjacent candidate transitions in parallel, and tree construction uses these precomputed scores to build a verification tree. Right: Tree-native Training. Our training combines high-temperature target traces with renormalized KL supervision to encourage plausible alternative branches.

Instruct Models
Model	Method	Math						Code						Chat
		GSM8K		M-500		AIME		HE		MBPP		LCB		MT-B		Avg.
		Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$
Qwen3-4B $T_{\mathrm{infer}}=0$	Single-path
	DFlash (16)	5.95	8.57	5.76	8.34	4.45	6.34	4.28	6.08	4.03	5.82	3.83	5.39	2.69	4.60	4.43	6.45
	Domino $\dagger$ (17)	8.02	10.12	7.14	9.01	5.97	7.49	5.66	7.02	5.59	7.11	5.33	6.86	3.28	5.12	5.86	7.53

Table 1: Speedup and average acceptance length (

\tau

) on instruct and reasoning models. All

\tau

values include the target-produced bonus token. M-500, HE, LCB, and MT-B denote MATH-500, HumanEval, LiveCodeBench, and MT-Bench.

\dagger

values are paper-reported.

Model	Method	$C=1$	$C=8$	$C=16$	$C=32$
Qwen3-4B	Autoregressive baseline	253.5 / $1.00\times$	1404.9 / $1.00\times$	2208.3 / $1.00\times$	3307.4 / $1.00\times$
	DFlash	796.4 / $3.14\times$	2287.2 / $1.63\times$	3664.4 / $1.66\times$	6071.9 / $1.84\times$
	Domino	816.3 / $3.22\times$	2258.9 / $1.61\times$	3636.5 / $1.65\times$	6115.5 / $1.85\times$
	DDTree	844.6 / $3.33\times$	2390.8 / $1.70\times$	3669.8 / $1.66\times$	4810.4 / $1.45\times$
	DARTree (adaptive $B/W$ )	978.5 / $3.86\times$	2405.1 / $1.71\times$	3770.3 / $1.71\times$	5619.9 / $1.70\times$
	UBTree (adaptive $B/W$ )	1089.1 / $4.30\times$	2797.4 / $1.99\times$	4271.9 / $1.93\times$	6560.4 / $1.98\times$

Table 2: Concurrent-serving results on Qwen3-4B and Qwen3-8B under

T_{\mathrm{infer}}=1

. Each cell reports aggregate output tokens/s / speedup over the AR baseline row at the same concurrency.

Model	Method	Math						Code						Chat
		GSM8K		M-500		AIME		HE		MBPP		LCB		MT-B		Avg.
		Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$
Ling3-Flash	DSpark	5.14	6.09	4.93	5.81	4.22	4.96	5.37	6.32	4.50	5.93	4.18	4.98	3.19	3.74	4.50	5.40
	DFlash	4.85	5.73	4.83	5.58	4.17	4.80	5.11	6.00	4.39	5.55	3.97	4.63	3.11	3.58	4.35	5.12
	UBTree	5.83	7.37	5.87	7.18	5.43	6.60	5.94	7.42	5.00	6.85	5.22	6.54	4.24	5.27	5.36	6.75
Qwen3.8-27B	DSpark	5.20	6.74	4.76	6.07	4.41	5.65	5.18	6.85	4.26	5.79	3.47	4.53	3.00	3.82	4.32	5.64

Table 3: Speedup and acceptance length (

\tau

) on production-scale models. Each benchmark reports speedup and average acceptance length, and the Avg. columns are means over all benchmarks. M-500, HE, LCB, and MT-B denote MATH-500, HumanEval, LiveCodeBench, and MT-Bench.

	Math						Code						Chat
	GSM8K		M-500		AIME		HE		MBPP		LCB		MT-B		Avg.
	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$	Speedup	$\tau$
DFlash 2 selector $+$ Tree	7.58	11.08	7.42	10.87	6.13	8.98	5.88	8.48	5.64	8.44	5.40	7.77	3.60	6.15	5.95	8.82
UBTree selector $+$ Tree	8.16	12.11	7.83	11.67	6.79	10.06	6.73	9.61	6.21	9.54	6.04	8.78	4.26	7.23	6.57	9.86
$+$ $T_{\mathrm{train}}=1$	8.54	12.58	8.18	12.15	6.92	10.13	7.07	10.22	6.66	9.98	6.37	9.24	4.39	7.32	6.88	10.23

Table 4: Selector, training-temperature, and training-objective ablations on Qwen3-4B with thinking disabled. Top: speedup and acceptance length (

\tau

) at

T_{\mathrm{infer}}=0

. Middle: acceptance length across training and inference temperatures, with the same checkpoints evaluated in both panels. Bottom: training-objective ablation with the fixed architecture and training temperature.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Value
Proposer training	6 epochs
Joint proposer–selector training	2 epochs
Selector rank	256
Codebook MLP hidden width	1024
Training candidate support size	256
Draft depth	15 future positions plus an anchor

Appendix

Table 5: UBTree configuration for the academic-scale experiments, unless otherwise specified. The verification budget excludes the anchor.

Benchmark	Examples / dialogues
GSM8K	128
MATH-500	128
AIME	30
HumanEval	128
MBPP	128
LiveCodeBench	128

Appendix

Table 6: Evaluation subset sizes. MT-Bench contains two-turn dialogues; the other counts refer to individual examples.

Method	Proposer stage	Draft post-processing	Verification	Other	Total	$\tau$	ms/token
DARTree	3.51	3.36	20.94	1.64	29.47	12.50	2.41
UBTree	3.41	1.38	20.88	1.61	27.30	12.58	2.21

Appendix

Table 7: Drafting Latency profiling on Qwen3-4B. Stage and total latencies are in ms per verification round.

Prompt tokens	$\leq 32$ K	32–64K	64–96K	96–128K	128–192K	192–260K	Overall
Requests	109	71	47	69	73	32	401
DSpark	4.248	4.297	3.795	3.976	4.166	4.084	4.136
DFlash	4.020	4.046	3.578	3.678	3.944	3.919	3.898
UBTree	5.833	5.870	5.326	5.351	5.798	5.804	5.693

Appendix

Table 8: Acceptance length (

\tau

) on LongBench-v2 with Ling3-Flash, grouped by request-start prompt length. All methods use the same 401 requests. The Overall column reports results over the full evaluation subset. Bold marks the best acceptance length per column.

Accumulated tokens	$\leq 0$	0–4K	4–8K	8–16K	16–32K	32–64K	64–131K	Overall
Covered cases	30	30	28	27	23	13	2	30
DSpark	3.718	3.743	4.023	3.998	4.089	4.023	3.965	3.981
DFlash	3.442	3.524	3.795	3.709	3.836	3.805	3.770	3.737
UBTree	4.842	4.829	5.447	5.389	5.516	5.408	5.465	5.341

Appendix

Table 9: Acceptance length (

\tau

) on SWE-bench with Ling3-Flash, grouped by accumulated prompt tokens relative to the first request in each case. Covered cases gives the number of cases represented in each bin. The Overall column reports request-pooled acceptance length over all 1,482 requests from 30 cases. Bold marks the best acceptance length per column.

Parallel Speculative Decoding

See all 28 papers

TreeGraft: Adaptive Multi-Drafter Grafting for Tree-Based Speculative Decoding
May 28, 2026Jiaming Fan, Daming Cao, Canchen Huang +6Speculative Decoding LLM Inference Acceleration
Key Laboratory of New Generation Artificial Intelligence Technology and Its Interdisciplinary Applications (Southeast University), Ministry of Education, China · Southeast University · Ant Group +1
TreeSpark: Calibrated, Load-Adaptive Draft Trees for Semi-Autoregressive Speculative Decoding
Aug 12, 2026Huapeng Zhou, Huayu Wang, Xinyu WangSpeculative Decoding LLM Inference Acceleration
Boson AI · University of Washington · McGill University
Unlocking Parallelism in Autoregressive Language Models via Speculative Decoding with Progressive Tree Drafting
Jul 12, 2026Zipeng Gao, Zhi Zheng, Qingrong Xia +5Speculative Decoding LLM Inference Acceleration
1State Key Laboratory of Cognitive Intelligence, University of Science and Technology of China · 2Huawei Technologies Co., Ltd.

May 28, 2026J/K move · Enter open · S save

TreeGraft: Adaptive Multi-Drafter Grafting for Tree-Based Speculative Decoding

Jiaming Fan, Daming Cao, Canchen Huang, Jiale Fu+5

Key Laboratory of New Generation Artificial Intelligence Technology and Its Interdisciplinary Applications (Southeast University), Ministry of Education, China · Southeast University · Ant Group+1

Speculative Decoding LLM Inference Acceleration LLM Inference Scheduling

Open paper AskPDF arXiv

Speculative decoding accelerates large language model inference through a draft-then-verify paradigm. Building on this, tree-structured methods improve inference by organizing proposals into multiple candidate paths, increasing the accepted length. However, existing tree-structured methods use a single drafter for all drafting steps, creating a dilemma: a smaller drafter is fast but yields lower-quality trees, whereas a larger drafter improves tree quality but suffers from high latency. To address this, we propose TreeGraft, a multi-drafter framework in which drafters of different costs jointly construct a shared draft tree. TreeGraft uses the stronger drafter to rescore candidates by updating scores assigned by the weaker drafter, reselect grafting positions, and recover promising paths left unexplored. It also integrates stronger drafter expansions non-destructively, preserving existing branches that may still be accepted by the target model. Together, these designs improve the quality of the shared draft tree. To control the drafting cost, TreeGraft introduces a lightweight scheduler distilled from an offline value system to decide when to call the stronger drafter. Across 10 model pairs and 6 benchmarks, TreeGraft outperforms the better of the two fixed single-drafter endpoint strategies by 15.1% on average, reaching a maximum gain of 26.6%. Our code is available at https://github.com/fjm9933/TreeGraft.

Abstract

Parallel Speculative Decoding

Explore similar work

TreeGraft: Adaptive Multi-Drafter Grafting for Tree-Based Speculative Decoding

TreeSpark: Calibrated, Load-Adaptive Draft Trees for Semi-Autoregressive Speculative Decoding

Unlocking Parallelism in Autoregressive Language Models via Speculative Decoding with Progressive Tree Drafting