PaperlayerView exact revision on arXiv

Source-linked AI summary

TreeSpark: Learning to Draft the Trees You Verify for Speculative Decoding

Jihao Zhao, Qianyu Jiang, Shuaishuai Zu, Pengfei Li, Guanghuan Fang, Ben Feng, Biao Qin, Zhibin Wang, Baochuan Yang

arXiv:2610.20175v1cs.CL

TL;DR

Parallel drafters can lose much of a proposed block after early mismatches, while reference-conditioned training does not directly supervise parent states and candidate choices introduced by retaining alternatives. TreeSpark replays deployed tree selection during training and adds survival and coverage objectives using reference target distributions. Across targets and suites, it improves mean acceptance length over DSpark by 23.4% to 31.6%.

  • Problem

    Reference-token losses do not directly supervise model-generated parent states or the candidate allocation and pruning decisions that determine which continuations reach verification.

    Source: §1 · pp. 1–2

  • Method

    TreeSpark replays deployed expansion, pruning, and packing on model-generated parents and combines path-mass-weighted survival with capped coverage using reused reference target distributions.

    Source: §1 · p. 2

  • Results

    Mean acceptance length over DSpark improves by 23.4% to 31.6% across targets and suites, with TreeSpark highest on all nine general benchmarks for both targets.

    Source: §4.2 · p. 7

  • Takeaways & Limitations

    Controlled comparisons support treating candidate admission and prefix retention as integral to draft learning alongside reference prediction accuracy.

    Source: §5 · p. 9

    Reference target distributions are exact for reference contexts and serve as proxies on alternative branches, without additional target evaluations.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

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.

1 INTRODUCTION

TreeSpark addresses the gap between fast parallel drafting and useful accepted continuations: reference-conditioned training does not directly supervise alternatives, model-generated parents, or candidate selection under a deployment budget. It aligns training with deployed tree selection while reusing reference target distributions.

1 INTRODUCTION

  • A single continuation loses its remaining suffix at the first rejection, while choosing only one path excludes alternatives represented in the drafter’s distributions.DSpark’s correction head is trained on reference parent tokens but encounters model-generated parents at deployment.Source: §1 · p. 1
  • Budgeted trees retain alternatives, but reference-token losses do not directly supervise which children receive verification positions or whether useful candidates survive pruning.Training must account for local prediction quality and the selection decisions that determine complete-path availability.Source: §1 · p. 2
  • Prior draft-generated data and online distillation align training with inference-time continuations and query distributions, while TreeSpark focuses on alignment for multiple retained continuations.Its correction head learns on proposed parents, reinforces retained paths, and promotes useful alternatives under the deployed budget.Source: §1 · p. 2
  • TreeSpark replays deployed expansion, pruning, and packing on model-generated parents, while retaining reference distillation and reusing reference target distributions.Survival weights retained-edge likelihood by reference-derived path mass; capped coverage promotes target-important alternatives near local selection boundaries.Source: §1 · p. 2
  • Its unified extension retains alternative continuations with one backbone pass and trains on model-generated parent states under the deployed candidate policy.Complementary survival and coverage objectives reuse reference target distributions without additional target forward passes.Source: §1 · p. 2

2 RELATED WORK

TreeSpark builds on parallel and semi-autoregressive drafting, inference-aligned training, and multi-candidate verification. Its focus is replaying deployment selections so supervision covers both retained paths and alternatives near candidate-admission boundaries.

2 RELATED WORK

  • Parallel drafting spans blockwise prediction, autoregressive target-feature drafting, DFlash’s block diffusion, and DSpark’s sequential token correction.TreeSpark builds on the DSpark Markov-head architecture and couples alternative-continuation generation with supervision on deployed parent states and candidate selections.Source: §2 · p. 2
  • Earlier alignment methods use draft-generated data, verification feedback, simulated rollouts, inference simulation, or acceptance-focused losses; TreeSpark adds alignment to branching candidate decisions.Its survival objective supervises retained paths, while coverage addresses alternatives near the admission boundary using reference target distributions.Source: §2 · pp. 2–3
  • Figure 1 depicts budgeted drafting and verification, coverage at expanded parents, and survival on packed edges weighted by reference path mass.Both objectives reuse reference target distributions; the schematic uses k = 2, while experiments use k = 4.Source: Figure 1 · p. 3
  • TreeSpark extends DSpark’s semi-autoregressive architecture by replaying deployment selections and coupling retained-path likelihood supervision with local admission signals for alternatives.The related work distinguishes this from candidate-tree construction and scoring approaches, including DDTree, TAPS, and JetSpec.Source: §2 · p. 3

3 ALIGNING PARALLEL DRAFTING WITH VERIFICATION

TreeSpark extends semi-autoregressive drafting by generating parent-conditioned candidate trees and replaying the deployed selection policy during training. Survival and coverage supervise different stages of selection, using reference target distributions while retaining shared backbone computation.

3.1 PARENT-CONDITIONED TREE DRAFTING

  • A single backbone pass supplies depth-specific logits, and lightweight token-dependent corrections produce parent-conditioned distributions without repeating the backbone.Each frontier parent proposes top-k children; depthwise frontier pruning bounds expansion, then global packing selects up to N nodes for verification.Source: §3.1 · pp. 3–4
  • Tree attention restricts each node to the committed prefix and its ancestor path; greedy verification follows matching target argmax children and emits a bonus token when that path ends.Under identical per-path target evaluations and argmax tie rules, the procedure preserves autoregressive greedy output for any such tree.Source: §3.1 · p. 4

3.2 TRAINING UNDER THE DEPLOYMENT POLICY

  • Training replays the current drafter’s tree policy on reference anchors, exposing correction heads to model-selected parent tokens while holding discrete selection decisions fixed.Reference target distributions from the aligned reference chain are reused across parents as proxies on alternative branches, avoiding additional target evaluations.Source: §3.2 · p. 4

3.3 SURVIVAL AND COVERAGE SUPERVISION

  • Survival weights packed-edge likelihood by reference-derived path mass, reducing the influence of descendants reached through low-mass prefixes.The packing restriction focuses likelihood supervision on candidates occupying verification slots; the objective is a surrogate for greedy acceptance.Source: §3.3 · pp. 4–5
  • Capped coverage penalizes target-favored tokens that cannot compete for local candidate slots, concentrating gradients within a bounded neighborhood of the selection threshold.It supervises expanded parents before global packing, including omitted alternatives within the active margin; survival instead supervises edges retained after packing.Source: §3.3 · p. 5

3.4 TRAINING AND INFERENCE

  • The complete objective adds survival and coverage losses to reference-token cross-entropy, distributional L1 matching, and the inherited confidence loss.The configuration uses B = 7, k = 4, and N = 32; tree objectives add no target evaluations during training or model components at inference.Source: §3.4 · p. 5
  • The draft attention head dimension is widened from 128 to 192 to increase capacity for learning shared-backbone branch decisions, with layer count and tree policy unchanged.The Markov correction rank remains unchanged, and the capacity choice is evaluated separately from the two objectives.Source: §3.4 · p. 5

4 EXPERIMENTS

TreeSpark improves acceptance and throughput over DSpark across the evaluated tasks and serving settings. The experiments also compare deployment-aligned tree supervision and attention capacity under controlled conditions.

4.1 EXPERIMENTAL SETUP

  • The evaluation uses greedy decoding with Qwen3-4B and Qwen3-8B targets, nine general benchmarks, and six writing domains.All drafters are trained for ten epochs on a shared mixture of 50,000 prompts; speculative methods allow up to seven draft depths, while TreeSpark packs 32 non-root nodes.Source: §4 · p. 5, §4.1 · pp. 5–6

4.3 SERVING EFFICIENCY AND TRAINING ALIGNMENT

  • Single-request throughput gains over DSpark average 15.1% and 17.9% on general tasks, and 21.9% and 25.0% on writing, for Qwen3-4B and Qwen3-8B.The throughput gains occur on every single-request benchmark.Source: Table 1 · p. 6, §4.3 · p. 8

4.2 ACCEPTANCE ACROSS TASKS AND DRAFT POSITIONS

  • Acceptance length rises over DSpark across both targets and task suites: 24.1% and 23.4% on general tasks, and 31.4% and 31.6% on writing.For Qwen3-4B and Qwen3-8B, general-suite means rise from 4.23 to 5.25 and from 4.20 to 5.18; writing means rise from 2.68 to 3.52 and from 2.55 to 3.35.Source: Table 2 · p. 7, §4.2 · p. 7
  • At C = 12, acceptance increases from 78.8% to 90.0% at draft position one and from 24.7% to 37.0% at position seven.The improvement persists at every plotted depth; Figure 2’s position-wise rates and recorded acceptance lengths come from separate evaluation records.Source: Figure 2 · p. 7, §4.2 · p. 7

4.3 SERVING EFFICIENCY AND TRAINING ALIGNMENT

  • At concurrent serving, Qwen3-4B system throughput gains over DSpark average 16.3% at C = 6, 17.2% at C = 12, and 12.6% at C = 36.Gains occur on all nine benchmarks at each tested concurrency.Source: Table 3 · p. 8, §4.3 · p. 8
  • With capacity and inference policy matched, TreeSpark raises mean acceptance from 4.93 to 5.25 and system throughput by 4.2% at C = 12 and 5.0% at C = 36.The control, DSpark + Tree, disables both tree losses; throughput means rise from 283.54 to 295.34 tokens/s at C = 12 and from 615.68 to 646.44 tokens/s at C = 36.Source: §4.3 · p. 8, Figure 3 · p. 9

4.4 CONTRIBUTIONS OF SUPERVISION AND CAPACITY

  • At C = 12, removing survival lowers the mean throughput gain over DSpark from 17.2% to 11.9%, while removing coverage lowers it to 16.0%.Survival has the larger aggregate effect; coverage adds a smaller improvement consistent with its role near local selection boundaries, though individual benchmarks sometimes favor an ablation.Source: §4.4 · pp. 8–9, Table 4 · p. 9
  • Increasing head dimension from 128 to 192 raises mean system throughput from 289.64 to 295.34 tokens/s, a 2.0% improvement.The 128-dimensional variant still improves over DSpark by 14.6%; the full configuration achieves the highest mean throughput in both ablation tables.Source: §4.4 · p. 9, Table 4 · p. 9

5 CONCLUSION

TreeSpark treats deployment policy as a guide for parallel-drafter training, directing reference supervision toward continuations available for target verification. Its appendix specifies how candidate trees are built and verified while preserving greedy target output under stated conditions.

5 CONCLUSION

  • TreeSpark’s controlled comparison yields longer mean accepted paths and higher mean throughput at fixed attention capacity and inference policy.The conclusion argues that candidate admission and prefix retention belong in draft learning alongside reference prediction accuracy.Source: §5 · p. 9

A TREE CONSTRUCTION AND GREEDY VERIFICATION

  • The construction retains alternative continuations from shared backbone outputs, and its verification arguments state conditions for preserving greedy target output.The guarantee assumes tree-masked and sequential target evaluations agree on identical paths and use the same argmax tie rule.Source: A TREE CONSTRUCTION AND GREEDY… · p. 13, §A.2 · p. 14

A.1 CANDIDATE EXPANSION AND GLOBAL PACKING

  • Candidate expansion retains children from shared backbone outputs, then globally packs up to N nodes by decreasing score, breaking ties by increasing depth.Nodes outside the next frontier remain eligible for verification; the resulting selected tree is prefix closed.Source: §A.1 · p. 13

A.2 ATTENTION CONTEXTS AND VERIFICATION

  • Tree-masked attention gives each node the committed prefix and its root-to-node path, while masking siblings and unrelated branches.Under matching path evaluations and argmax tie rules, the verifier emits the same tokens as autoregressive greedy decoding.Source: §A.2 · p. 14

A.3 ACCEPTANCE LENGTH CONVENTION

  • Acceptance length counts accepted non-root draft nodes plus the target bonus, excluding the already-committed anchor.Rounds ending at an accepted end token or output limit can omit the bonus.Source: §A.3 · p. 14

B TRAINING OBJECTIVES AND REPLAY DETAILS

Training replay exposes the drafter to its proposed parent states while reusing reference supervision. Survival weights packed edges by reference path mass, and coverage promotes target-important alternatives; baseline losses and distributed normalization remain part of the objective.

B.1 REFERENCE SUPERVISION AND REPLAY

  • Reference-derived path mass recursively multiplies each parent’s mass by the reference probability of the proposed child token.Later-depth reference distributions are shared across parents even when their paths differ from the reference.Source: §B.1 · p. 15

B TRAINING OBJECTIVES AND REPLAY DETAILS

  • Replay trains on proposed parent states under the current tree policy, combining survival and coverage losses with baseline reference objectives.Candidate identities and joint scores are detached, while the topology is rebuilt after parameter updates.Source: B TRAINING OBJECTIVES AND REPLAY DETAILS · p. 14, §B.1 · p. 15

B.2 PACKED-EDGE SURVIVAL OBJECTIVE

  • Packed-edge survival weights retained-child log probabilities by detached reference path mass and the packed-node mask, without explicit depth decay.Conditional log probabilities remain differentiable.Source: §B.2 · p. 15

B.3 FRONTIER COVERAGE OBJECTIVE

  • Frontier coverage weights parent penalties by target-token probability and depth, giving earlier candidate decisions greater weight.Expanded parents contribute even when their children are later excluded by packing.Source: §B.3 · pp. 15–16

B.4 BASELINE LOSSES AND DISTRIBUTED NORMALIZATION

  • Baseline losses use masked, depth-decayed reductions, while distributed normalization divides local numerators by globally summed valid weights.The inherited confidence objective uses overlap as one-step acceptance probability for ordinary rejection sampling, not as a substitute for greedy tree acceptance depth.Source: §B.4 · p. 16

C ANALYSIS OF THE SUPERVISION OBJECTIVES

The analysis connects reference-derived path mass to expected tree coverage and explains how survival and coverage produce complementary training signals. It also distinguishes packed-tree diagnostics from deployed greedy acceptance.

C.1 PATH MASS AS A SUPERVISION WEIGHT

  • The depthwise coverage identity expresses expected contained-prefix length through the probability that sampled target prefixes occur in the tree.At each depth, distinct nodes represent disjoint prefixes, so the probability of a contained prefix is the sum of their path probabilities; the identity motivates detached survival weights during replay, while the implemented objective remains a supervised surrogate.Source: §C.1 · pp. 16–17
  • Reference-derived path masses measure coverage under an auxiliary product distribution, approximating target coverage when reference distributions represent alternative paths.Coverage is nonincreasing with depth in a prefix-closed tree because each extension adds a constraint.Source: §C.1 · p. 17

C.2 SURVIVAL GRADIENTS ON RETAINED EDGES

  • Survival assigns retained children path-mass-weighted likelihood targets, while parent path mass scales each parent's contribution across states.Omitted tokens receive only normalization gradients from survival, motivating coverage and reference distillation.Source: §C.2 · p. 17

C.3 COVERAGE GRADIENTS AND SATURATION

  • Coverage promotes target-important tokens near the selection boundary, but its capped hinge gives no direct gradient to tokens far below it.The loss is a local margin surrogate rather than an exact set-membership loss, and it can remain nonzero when the target top-J set is already included in the draft top-k set. Its normalized value is bounded by c.Source: §C.3 · p. 17

C.4 DIAGNOSTICS ON THE PACKED TREE

  • Packed-tree diagnostics measure reference-mass coverage and reference-argmax survival, whereas deployment evaluation measures recursively conditioned greedy acceptance.The reference-argmax and greedy sequences coincide when their conditioning contexts agree; batch depth averages can have different denominators.Source: §C.4 · p. 18

D IMPLEMENTATION AND COMPUTATIONAL COST

Implementation keeps tree branching on a shared backbone and lightweight correction head, while matching training and serving selection policies. The section details the configured tree budget, replay costs, batching, and numerical selection conventions.

D.1 CONFIGURATION AND CAPACITY

  • The reference configuration uses B = 7, k = 4, and N = 32, with up to 100 candidates before packing and 33 target positions including the anchor.The correction rank is 256, and deployment requires matching the recorded tree hyperparameters and inference arguments.Source: §D.1 · p. 18
  • Draft attention head dimension 192 localizes the capacity increase to draft attention, while the residual hidden size, layer count, Markov rank, and tree budget remain unchanged.The affected projection sizes and draft KV storage scale by 1.5, not the entire model's parameter count or memory footprint.Source: §D.1 · p. 18

D.2 REPLAY COST AND ACTIVATION MEMORY

  • Shared backbone outputs serve all parents, while replay computes at most k rank-r Markov corrections per anchor and depth without repeating transformer computation for branches.Tree supervision adds draft computation and backward work, but reuses teacher information already needed for reference distillation; checkpointing reconstructs depth logits to limit retained activations.Source: §D.2 · pp. 18–19

D.3 BATCHED CONSTRUCTION AND VERIFICATION

  • Batched frontier expansion and verification share depth-wise correction calls across requests, and packed target trees allow requests to commit different token counts.Training and serving packed sets coincide for matching scores without boundary ties; serving's depth-aware tie rule differs from training's floating-point top-N selection.Source: §D.3 · p. 19, §D.4 · p. 19

D.4 NUMERICAL SELECTION CONVENTIONS

  • Joint path comparison requires normalized row scores, while serving computes selected-token values without materializing a full normalized vocabulary tensor.Serving's input-precision selection can differ from full float32 log-softmax near selection boundaries; the stated target-verification conditions preserve greedy output.Source: §D.3 · p. 19, §D.4 · p. 19

E.1 DATA AND TRAINING PROTOCOL

The evaluation uses a common 50,000-prompt training mixture and evaluates decoding efficiency on general and writing workloads. Comparisons retain common training settings and distinguish the reference-trained tree control from capacity and loss ablations.

E.1 DATA AND TRAINING PROTOCOL

  • The common 50,000-prompt training mixture combines mathematics, code, dialogue and instruction, and writing prompts, with writing sourced from Writer-R1.Its final composition is 29.92% mathematics, 40.15% code, 19.94% dialogue and instructions, and 9.99% writing; training sequences are capped at 4,096 tokens.Source: §E.1 · pp. 19–20, Table 5 · p. 20
  • The general evaluation suite covers nine mathematics, coding, and dialogue or instruction benchmarks, while WritingBench covers six writing domains.The reported values measure throughput and acceptance, not task accuracy or writing quality; all compared methods use the same selected general-suite prompts.Source: §E.1 · p. 20
  • All drafters use ten epochs, AdamW, a 6 × 10^-4 learning rate, bfloat16 computation, and the final checkpoint for results.The DSpark + Tree control shares TreeSpark's head dimension and inference policy but disables both tree losses; single-loss ablations disable one coefficient, and the capacity ablation uses dimension 128.Source: §E.1 · pp. 20–21

E.2 METRIC DEFINITIONS AND AGGREGATION

The section defines acceptance and throughput measurements and explains how benchmark means and relative gains are aggregated. It also reports writing-suite acceptance results and notes that throughput uncertainty is not quantified.

E.2 METRIC DEFINITIONS AND AGGREGATION

  • Acceptance length counts newly committed tokens per verification round, including the target bonus when emitted, while system throughput counts generated tokens per elapsed workload time.Single-request throughput divides output tokens by accumulated decoding time; per-request throughput excludes time to first token and need not equal system throughput divided by concurrency.Source: §E.2 · p. 21
  • Mean benchmark metrics are unweighted, but throughput-table ∆ means average benchmark-level relative gains rather than the relative change between mean throughput values.Figure 3 and the stated tree-configuration comparisons use arithmetic mean absolute metrics.Source: §E.2 · p. 21
  • Throughput values average three repeated measurements, but the report provides no error bars or significance tests and does not quantify repetition uncertainty.Table 7 presents per-request throughput on Qwen3-4B with identical hardware and concurrency across methods.Source: §E.2 · p. 21, Table 7 · p. 22

E.3 ACCEPTANCE ON WRITINGBENCH

  • TreeSpark has the highest acceptance length in every WritingBench domain for both targets, with gains over DSpark of 30.3%–32.4% for Qwen3-4B and 29.5%–32.6% for Qwen3-8B.Table 6 reports acceptance length across the six domains, with the mean weighting domains equally.Source: §E.3 · p. 21, Table 6 · p. 22

E.4 PER-REQUEST THROUGHPUT AND ABLATIONS

Per-request throughput increases across the reported concurrency settings and benchmarks. Ablations show smaller gains when either training objective is removed or the attention head is narrower.

E.4 PER-REQUEST THROUGHPUT AND ABLATIONS

  • 29.39 to 33.63 tokens/s/request at C = 6, 23.76 to 27.86 at C = 12, and 22.08 to 24.55 at C = 36; every benchmark improves.Table 8 reports per-request throughput, not tail latency.Source: §E.4 · p. 22, Table 8 · p. 23
  • The 17.6% average per-request throughput gain over DSpark falls to 12.8% without survival, 15.0% without coverage, and 15.3% with head dimension 128.Each variant retains tree inference.Source: §E.4 · p. 22, Table 8 · p. 23
  • At fixed objectives and tree policy, widening the head from 128 to 192 raises mean per-request throughput from 27.29 to 27.86 tokens/s, a 2.1% gain, and mean system throughput by 2.0%.The 2.1% gain is computed before rounding.Source: §E.4 · p. 22
  • The full configuration has the highest per-request throughput on each benchmark; together with Table 4, the results support combining tree supervision with the larger attention representation.Table 8 marks column maxima and compares variants with DSpark.Source: §E.4 · p. 23, Table 8 · p. 23

E.5 TRAINING ALIGNMENT ACROSS BENCHMARKS

TreeSpark and DSpark + Tree exceed DSpark in mean throughput at the reported concurrency settings, while aligned supervision adds gains at fixed capacity and inference policy. The evaluation is limited to greedy decoding with a fixed candidate budget, and the manuscript discloses AI assistance in coding and writing.

E.5 TRAINING ALIGNMENT ACROSS BENCHMARKS

  • Figures 4–6 enlarge the same system and per-request measurements rather than adding experimental runs.Figure 4 compares DSpark, DSpark + Tree, and TreeSpark at total concurrency 12 and 36; Figures 5 and 6 separate system and per-request views by benchmark.Source: §E.5 · p. 23, Figure 4 · p. 24, Figure 5 · p. 25
  • At C = 12, mean system throughput is 254.68 tokens/s for DSpark, 283.54 for DSpark + Tree, and 295.34 for TreeSpark; at C = 36, it is 574.45, 615.68, and 646.44, respectively.At fixed capacity and inference policy, the two tree configurations have mean acceptance lengths of 4.93 and 5.25.Source: §E.5 · p. 23, Figure 4 · p. 24
  • The reference-trained control leads system throughput on AIME 2025, MBPP, and ArenaHard-v2 at C = 12, and on GSM8K, MT-Bench, and ArenaHard-v2 at C = 36, while the complete system exceeds chain DSpark throughout.The passage distinguishes the complete-system comparison from the fixed-capacity supervision comparison.Source: §E.5 · p. 23
  • The evaluation uses two Qwen3 target sizes, greedy decoding, and a fixed candidate budget; generalization across budgets and stochastic sampling remains unevaluated.Greedy output preservation is established under the verification assumptions in Appendix A.Source: §E.5 · p. 23

F AI ASSISTANCE DISCLOSURE

  • AI assistance supported parts of software development and academic writing, with assisted code and text reviewed by the authors.The authors state AI tools did not formulate research questions, determine methodology, interpret results, or derive conclusions.Source: F AI ASSISTANCE DISCLOSURE · pp. 23–24

Related papers

Loading 2610.20175v1…