TreeSpark: Learning to Draft the Trees You Verify for Speculative Decoding
Authors: Jihao Zhao, Qianyu Jiang, Shuaishuai Zu, Pengfei Li, Guanghuan Fang, Ben Feng, Biao Qin, Zhibin Wang, Baochuan Yang
Organizations: School of Information, Renmin University of China · National Key Lab for Novel Software Technology, Nanjing University · Huawei Technologies Ltd.
Parallel drafters such as DFlash and DSpark reduce proposal latency, but early mismatches limit how much of each block can be accepted. Retaining alternatives introduces parent states and candidate choices that reference-conditioned training does not directly supervise. We introduce TreeSpark, a unified training and inference framework extending the semi-autoregressive DSpark architecture. It constructs parent-conditioned candidate trees from one backbone pass and replays their expansion, pruning, and budgeted packing during training. Survival negative log-likelihood weights retained edges by reference-derived path mass, while a capped coverage hinge promotes target-important alternatives near local selection boundaries. Both objectives reuse reference-chain target distributions without additional target forward passes. We evaluate greedy decoding with Qwen3-4B and Qwen3-8B on nine general benchmarks and six writing domains. Across targets and suites, TreeSpark improves mean acceptance length over DSpark by 23.4% to 31.6% and mean single-request throughput by 15.1% to 25.0%. Batched Qwen3-4B throughput improves by 12.6% to 17.2%. At fixed attention capacity and inference policy, the proposed supervision further improves mean acceptance length by approximately 6.4% and mean system throughput by 4.2% to 5.0% on Qwen3-4B.
Figures & tables
Figure 1: TreeSpark framework. (a) Budgeted drafting and verification. (b) Coverage at expanded parents. (c) Survival on valid packed edges, weighted by reference path mass. Both objectives reuse reference target distributions. Teal nodes are packed; double outlines mark the frontier. The schematic uses k=2 ; experiments use k=4 . Panel (c) illustrates one supervised path. Equation 1 defines parent conditioning. The feedback loop updates draft parameters; global packing follows depthwise expansion.
(a) Reasoning, code, and dialogue
GSM8K
MATH 500
AIME 25
Human Eval
MBPP
LiveCode Bench
MT- Bench
Alpaca
Arena Hard-v2
Mean
Qwen3-4B
Vanilla
12.83
11.66
10.82
12.80
13.13
11.53
12.99
12.78
11.72
12.25
EAGLE-3
25.20
23.27
21.29
20.75
21.19
22.89
16.06
15.51
16.03
20.24
DFlash
46.66
43.48
37.48
35.36
36.59
42.88
24.21
22.87
23.58
34.79
DSpark
49.74
47.29
41.22
39.39
37.22
42.06
24.30
23.06
24.05
36.48
Table 1: Single-request throughput (tokens/s): (a) reasoning, code, and dialogue; (b) writing. Vanilla denotes autoregressive decoding. Mean averages benchmarks. Δ compares each benchmark with DSpark; its Mean entry averages these gains. Bold marks column maxima.
Reasoning, code, and dialogue
GSM8K
MATH 500
AIME 25
Human Eval
MBPP
LiveCode Bench
MT- Bench
Alpaca
Arena Hard-v2
Mean
Qwen3-4B
EAGLE-3
3.43
3.25
3.04
2.83
2.76
3.20
2.13
2.08
2.20
2.77
DFlash
4.97
4.95
4.42
3.91
3.76
4.82
2.52
2.41
2.59
3.82
DSpark
5.46
5.44
4.97
4.41
4.22
5.24
2.80
2.63
2.90
4.23
TreeSpark
6.57
6.58
6.26
5.54
5.40
6.23
3.59
3.42
3.64
5.25
Table 2: Acceptance length on reasoning, code, and dialogue tasks, including the target bonus when emitted. All methods allow seven draft depths; TreeSpark packs 32 non-root nodes. Mean averages benchmarks; its Δ compares the two means for TreeSpark and DSpark. Other Δ entries compare individual benchmarks. Bold marks column maxima. Writing results appear in Table 6 .
Figure 2: Acceptance on Qwen3-4B. Left: position-wise acceptance at C=12 , averaged over nine benchmarks; shading spans their minimum and maximum. Right: recorded acceptance lengths from Table 2 , including the target bonus when emitted. The panels use separate evaluation records.
System throughput (tokens/s)
GSM8K
MATH 500
AIME 25
Human Eval
MBPP
LiveCode Bench
MT- Bench
Alpaca
Arena Hard-v2
Mean
C=6
Vanilla
0 68.79
0 58.50
0 45.27
0 59.66
0 62.80
0 49.00
0 56.86
0 59.37
0 54.71
0 57.22
DSpark
235.51
223.87
183.69
173.60
167.44
181.85
0 92.84
106.23
118.01
164.78
TreeSpark
256.65
256.29
203.90
201.66
179.66
208.52
127.50
124.15
141.94
188.92
Δ (%)
+9.0
+14.5
+11.0
+16.2
+7.3
+14.7
+37.3
+16.9
+20.3
+16.3
Table 3: System throughput on Qwen3-4B (tokens/s) at total concurrency C . Serving configurations are specified in Section 4.1 . Δ is the gain over DSpark within each setting; its Mean entry averages per-benchmark gains.
Figure 3: Effect of training alignment on Qwen3-4B. DSpark + Tree matches the capacity and inference policy of TreeSpark but omits both tree losses. All panels report means over nine benchmarks. DSpark and TreeSpark acceptance means match Table 2 .
System throughput (tokens/s)
GSM8K
MATH 500
AIME 25
Human Eval
MBPP
LiveCode Bench
MT- Bench
Alpaca
Arena Hard-v2
Mean
DSpark
341.53
374.50
274.46
258.41
250.87
284.77
139.06
173.52
195.02
254.68
TreeSpark
381.60
435.55
314.90
308.42
281.77
316.64
186.24
209.82
223.15
295.34
Δ (%)
+11.7
+16.3
+14.7
+19.4
+12.3
+11.2
+33.9
+20.9
+14.4
+17.2
w/o Lsurv
364.20
412.68
297.14
297.21
281.76
292.01
171.53
200.31
220.20
281.89
Δ (%)
0 +6.6
+10.2
0 +8.3
+15.0
+12.3
0 +2.5
+23.3
+15.4
+12.9
+11.9
Table 4: Component ablations on Qwen3-4B at C=12 (system tokens/s). Loss ablations remove only the indicated objective; the capacity ablation uses head dimension 128 instead of 192. Δ compares each variant with DSpark; its Mean entry averages benchmark gains. Bold marks column maxima.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Category
Source
Records
Share (%)
Mathematics
GSM8K
7,450
14.90
MATH
7,480
14.96
AIME 2024
30
0.06
Code
MBPP
120
0.24
Magicoder OSS-Instruct, Python
9,481
18.96
APPS
3,989
7.98
Appendix
Table 5: Composition of the 50,000-prompt training mixture after length filtering, sampling, shuffling, and truncation. Counts refer to records, not unique prompt strings.
Writing
Academic & Engineering
Advertising & Marketing
Education
Finance & Business
Literature & Arts
Politics & Law
Mean
Qwen3-4B
EAGLE-3
2.09
1.98
2.06
2.19
1.94
2.09
2.06
DFlash
2.51
2.31
2.46
2.62
2.22
2.47
2.43
DSpark
2.79
2.51
2.72
2.90
2.41
2.75
2.68
TreeSpark
3.68
3.27
3.58
3.84
3.16
3.60
3.52
Appendix
Table 6: Acceptance length on the six WritingBench domains, supplementing Table 2 . The target bonus is included when emitted. Mean weights domains equally; its Δ compares the two means for TreeSpark and DSpark. Other Δ entries compare individual domains. Bold marks column maxima.
Per-request throughput (tokens/s/req)
GSM8K
MATH 500
AIME 25
Human Eval
MBPP
LiveCode Bench
MT- Bench
Alpaca
Arena Hard-v2
Mean
C=6
Vanilla
11.69
0 9.89
0 8.73
10.93
10.92
0 9.05
10.24
10.38
0 9.30
10.13
DSpark
40.44
38.58
34.25
31.59
29.51
32.56
18.61
18.51
20.46
29.39
TreeSpark
44.09
44.42
38.24
36.13
33.84
36.36
23.10
22.05
24.46
33.63
Δ (%)
+9.0
+15.1
+11.6
+14.4
+14.7
+11.7
+24.1
+19.1
+19.6
+15.5
Appendix
Table 7: Per-request decoding throughput on Qwen3-4B (tokens/s/request), supplementing Table 3 . Hardware and concurrency are identical. Δ gives relative gains over DSpark; its Mean entry averages per-benchmark gains.
Per-request throughput (tokens/s/req)
GSM8K
MATH 500
AIME 25
Human Eval
MBPP
LiveCode Bench
MT- Bench
Alpaca
Arena Hard-v2
Mean
DSpark
29.52
33.21
29.94
24.81
22.60
25.15
15.85
15.42
17.38
23.76
TreeSpark
32.99
38.44
36.76
29.50
26.24
28.52
19.73
18.66
19.87
27.86
Δ (%)
+11.8
+15.7
+22.8
+18.9
+16.1
+13.4
+24.5
+21.0
+14.3
+17.6
w/o Lsurv
31.54
36.88
35.30
28.61
25.20
26.72
18.43
18.11
19.52
26.70
Δ (%)
0 +6.8
+11.1
+17.9
+15.3
+11.5
0 +6.2
+16.3
+17.4
+12.3
+12.8
Appendix
Table 8: Per-request throughput for the ablations in Table 4 (tokens/s/request). Each variant retains tree inference. Δ compares each variant with DSpark; its Mean entry averages benchmark gains. Bold marks column maxima.
Figure 4: Detailed throughput on Qwen3-4B. System and per-request measurements at total concurrency 12 and 36 compare DSpark, DSpark + Tree, and TreeSpark. The two tree configurations use attention head dimension 192.
Figure 5: System throughput by benchmark. Enlarged view of the system measurements in Figure 4 under the serving configurations in Section 4.1 .
Figure 6: Per-request throughput by benchmark. Enlarged view of the per-request measurements in Figure 4 , averaged over request decoding intervals after the first output token.
Key Laboratory of New Generation Artificial Intelligence Technology and Its Interdisciplinary Applications (Southeast University), Ministry of Education, China · Southeast University · Ant Group +1