Robot2050 / TreeSpark

TreeSpark: Learning to Draft the Trees You Verify for Speculative Decoding
## 🎯 Who Should Pay Attention to Our Work?
* **LLM Inference and Serving Engineers:** Parallel drafters such as DFlash and DSpark lose the remainder of each block to the first rejection. TreeSpark turns the same backbone computation into a budgeted tree of alternatives, so more of every computed block survives verification.
* **Speculative Decoding Researchers:** Reference-conditioned training never supervises the parent states or the candidate selections that tree verification actually uses. TreeSpark closes this gap by replaying the deployed candidate policy during training.
* **Draft Distillation Practitioners:** TreeSpark provides inference-aligned supervision without placing a frozen target model in the training loop. Both tree objectives reuse the reference-chain target distributions already computed for distillation.
* **High-Concurrency Serving Developers:** The bundled continuous-batching engine packs and verifies all in-flight draft trees in a single target forward per round, with device-side mask assembly and vectorized acceptance retrieval.
## ✨ Core Contributions
1. **Parent-Conditioned Tree Drafting from One Backbone Pass.**
 One draft backbone forward produces vocabulary logits for every draft depth, and a lightweight first-order Markov correction head specializes each depth to a specific parent token (`z_d(u) = ℓ_d(x) + W·e(t_u)`). Branching therefore requires only extra head evaluations and never repeats the backbone. Trees expand with branching factor `k=4` over `B=7` depths and are packed under a global budget of `N=32` non-root nodes: up to 100 candidates before packing, and 33 target input positions including the anchor.
2. **Training under the Deployment Policy.**
 Training replays the exact candidate policy used at serving time, namely per-parent top-k expansion, joint path-score frontier pruning, and global top-N packing, on model-generated parent states. Discrete selections (token identities, frontier choices, packing decisions) are detached, while the conditional distributions remain differentiable through the backbone and the Markov head. This removes the exposure mismatch between teacher-forced training and tree deployment without any online target rollouts.
3. **Complementary Survival and Coverage Objectives.**
 Survival NLL weights the likelihood of each packed edge by its reference-derived path mass `m(v)`, a differentiable surrogate for how deeply target probability mass survives inside the verified tree. The capped coverage hinge (`J=4`, margin `δ=1`, cap `c=4`, depth decay `γ=4`) promotes target-important tokens near the local top-k admission boundary, including candidates that would otherwise be omitted and receive no survival signal. The two terms act at complementary stages of selection: coverage supervises expanded parents before the node budget is applied, and survival supervises the edges retained after packing.
4. **Exactness by Con