pith. sign in

arxiv: 2606.03819 · v1 · pith:3ARFMTK5new · submitted 2026-06-02 · 💻 cs.LG

TreeFlash: Parallel AR-Approximation for Faster Speculative Decoding

Peer Rheinboldt , Fr\'ed\'eric Berdoz , Roger Wattenhofer This is my paper

Pith reviewed 2026-06-28 10:53 UTC · model grok-4.3

classification 💻 cs.LG
keywords speculative decodingtree draftingautoregressive approximationone-shot draftersblock efficiencyMLP layerconstant-time decoding
0
0 comments X

The pith

TreeFlash adds an MLP to one-shot drafters to approximate autoregressive token distributions while keeping constant-time tree drafting.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that one-shot block drafters for speculative decoding suffer from non-autoregressive conditioning, which makes their token predictions diverge from the verifier model's true distribution as depth increases, and this effect worsens in tree structures where branches must share the same marginal. TreeFlash inserts an MLP layer that takes the drafter hidden state plus the previous token to restore approximate autoregressive dependence. The mechanism uses a two-stage process so the entire draft tree can still be produced in one forward pass. A sympathetic reader would care because closer draft distributions raise acceptance rates, which directly reduces wasted computation in large-model inference.

Core claim

TreeFlash incorporates an MLP layer conditioned on the drafter's hidden state and the previous token to approximate an autoregressive distribution. It retains the O(1) decoding time complexity of one-shot drafters by employing a two-stage approximation mechanism. TreeFlash achieves state-of-the-art performance across a variety of tasks and models, improving over marginal tree drafting by 12% higher block efficiency and 9% higher speedup.

What carries the argument

The MLP layer conditioned only on the drafter hidden state and previous token, used inside a two-stage approximation that produces the full draft tree in constant time.

If this is right

  • Tree-based speculative decoding can now use deeper or wider trees without the previous divergence penalty.
  • Block efficiency rises because drafted tokens better match what the verifier would have chosen.
  • Overall generation speedup increases by roughly nine percent on the tested models and tasks.
  • The one-shot property is preserved, so throughput scaling with batch size remains unchanged.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same MLP trick could be added to other non-autoregressive drafting schemes that currently ignore token-to-token dependence.
  • If the approximation holds at greater depths, practitioners could safely increase draft-tree branching factors to raise acceptance rates further.
  • The approach might reduce the engineering effort spent on designing hand-crafted tree topologies.

Load-bearing premise

The MLP produces a sufficiently close approximation to the verifier's true autoregressive conditional distribution even as tree depth increases, without requiring sequential computation.

What would settle it

Run the same tree-drafting experiments but measure the KL divergence between TreeFlash's approximated next-token distribution and the verifier's true distribution at successive depths; if the divergence grows faster than in marginal drafting and block efficiency gains disappear, the central claim is false.

Figures

Figures reproduced from arXiv: 2606.03819 by Fr\'ed\'eric Berdoz, Peer Rheinboldt, Roger Wattenhofer.

Figure 1
Figure 1. Figure 1: Speedup averaged across datasets for stan [PITH_FULL_IMAGE:figures/full_fig_p001_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Overview of different drafting paradigms. EAGLE-3 ( [PITH_FULL_IMAGE:figures/full_fig_p002_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Total Variation Distance (TVD) to the verifier distribution across draft positions for Qwen3 4B. We compare DFlash (q(xt+i |x≤t)), TreeFlash (q ′ (xt+i |x≤t, xt+i−1)), and an approxi￾mation of the ground-truth marginal distribution p(xt+i |x≤t), which is generated using Monte Carlo sampling from the verifier distribution. Both DFlash and TreeFlash start with a relatively low TVD of 0.19 at the first token.… view at source ↗
Figure 4
Figure 4. Figure 4: Coverage of target probability space with [PITH_FULL_IMAGE:figures/full_fig_p005_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Block efficiency and throughput for different [PITH_FULL_IMAGE:figures/full_fig_p007_5.png] view at source ↗
Figure 7
Figure 7. Figure 7: Qualitative comparison of decoding sub-trees produced by DDTree and TreeFlash. Outline tokens indicate [PITH_FULL_IMAGE:figures/full_fig_p008_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: is the algorithm used for efficient tree con￾struction. Note that Part 1 can be done efficiently in parallel on the GPU and Part 2 is done iteratively on the CPU. C Qualitative Trees As illustrated in Figures 9 and 10, the two methods produce qualitatively different tree topologies. Be￾cause DDTree relies solely on the marginal distri￾bution, the token ranking at each depth is identical across all paths. T… view at source ↗
Figure 9
Figure 9. Figure 9: Example of a Draft Tree with B=64 produced by DDTree. Outlined tokens are accepted by Qwen3 8B under greedy sampling. 12 [PITH_FULL_IMAGE:figures/full_fig_p012_9.png] view at source ↗
Figure 10
Figure 10. Figure 10: Example of a Draft Tree with B=64 produced by TreeFlash. Outlined tokens are accepted by Qwen3 8B under greedy sampling. 13 [PITH_FULL_IMAGE:figures/full_fig_p013_10.png] view at source ↗
read the original abstract

One-shot block drafters for speculative decoding generate the full draft in a single forward pass, achieving strong throughput by eliminating sequential token generation. However, they predict each draft token conditioned only on the prefix context, with no dependence on previously drafted tokens. This non-autoregressive conditioning causes the drafter's distribution to diverge from the verifier's true autoregressive distribution as draft depth grows. This problem becomes more severe in tree-based drafting, where distinct branches are forced to share the same marginal distribution for subsequent tokens. We propose TreeFlash, which addresses this by incorporating an MLP layer conditioned on the drafter's hidden state and the previous token to approximate an autoregressive distribution. TreeFlash retains the $\mathcal{O}(1)$ decoding time complexity of one-shot drafters by employing a two-stage approximation mechanism. TreeFlash achieves state-of-the-art performance across a variety of tasks and models, improving over marginal tree drafting by $12\%$ higher block efficiency and $9\%$ higher speedup.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

3 major / 1 minor

Summary. The manuscript introduces TreeFlash, a method for one-shot block drafters in speculative decoding that adds an MLP layer conditioned on the drafter hidden state and previous token to approximate an autoregressive distribution. This is implemented via a two-stage mechanism claimed to preserve O(1) complexity, addressing divergence issues in tree-based drafting where branches share marginal distributions. The paper reports state-of-the-art results, including 12% higher block efficiency and 9% higher speedup over marginal tree drafting across tasks and models.

Significance. If the approximation fidelity holds without introducing hidden sequential costs, the work could improve throughput in speculative decoding for LLMs by better aligning drafter and verifier distributions in parallel tree settings. The empirical gains, if reproducible, would be a practical engineering advance in inference optimization.

major comments (3)
  1. [Abstract] Abstract: the performance claims of 12% higher block efficiency and 9% higher speedup are presented without any description of the training procedure, loss function, dataset details, or ablation results. This omission prevents verification that the gains arise from the MLP approximation rather than post-hoc selection or unstated baselines.
  2. [Abstract] Abstract: no equations, pseudocode, or complexity derivation is supplied for the two-stage approximation mechanism or the MLP conditioning on hidden state plus previous token. Without these, the claim that the method maintains O(1) decoding time while approximating the verifier's autoregressive conditional cannot be checked.
  3. [Abstract] Abstract: the central assumption that the MLP yields a distribution sufficiently close to the verifier's true autoregressive conditional as tree depth increases and branches diverge lacks any error-vs-depth measurements, receptive-field analysis, or branch-divergence experiments. This directly bears on whether the reported efficiency numbers can be attributed to the method.
minor comments (1)
  1. The abstract would be clearer if it briefly listed the specific models, tasks, and baseline implementations used to support the SOTA claim.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for the constructive comments. We address each major comment point by point below, clarifying where details appear in the manuscript and indicating revisions where appropriate.

read point-by-point responses

  1. Referee: [Abstract] Abstract: the performance claims of 12% higher block efficiency and 9% higher speedup are presented without any description of the training procedure, loss function, dataset details, or ablation results. This omission prevents verification that the gains arise from the MLP approximation rather than post-hoc selection or unstated baselines.

    Authors: The abstract is a concise summary. The training procedure (two-stage supervised fine-tuning), loss function (KL divergence to verifier autoregressive targets), dataset details (standard next-token prediction corpora), and ablation results isolating the MLP contribution are provided in Sections 3.2, 3.3, and 4.3. These sections confirm the gains derive from the approximation. We will revise the abstract to reference the training setup and ablations. revision: yes

  2. Referee: [Abstract] Abstract: no equations, pseudocode, or complexity derivation is supplied for the two-stage approximation mechanism or the MLP conditioning on hidden state plus previous token. Without these, the claim that the method maintains O(1) decoding time while approximating the verifier's autoregressive conditional cannot be checked.

    Authors: The abstract omits technical detail for brevity. The two-stage mechanism is defined in Equations (1)-(4) of Section 3.1, MLP conditioning on hidden state and prior token appears in Equation (2), pseudocode is given in Algorithm 1, and the O(1) complexity derivation (parallel forward pass plus constant-time MLP) is in Section 3.3. We will update the abstract to briefly describe the two-stage mechanism. revision: partial

  3. Referee: [Abstract] Abstract: the central assumption that the MLP yields a distribution sufficiently close to the verifier's true autoregressive conditional as tree depth increases and branches diverge lacks any error-vs-depth measurements, receptive-field analysis, or branch-divergence experiments. This directly bears on whether the reported efficiency numbers can be attributed to the method.

    Authors: Section 4.4 reports approximation error versus draft depth, demonstrating bounded divergence. Receptive-field analysis for the MLP appears in Section 3.2. Branch divergence is addressed via the tree-drafting efficiency comparisons in Section 4.2. These results support attribution of the reported gains to the method. revision: no

Circularity Check

0 steps flagged

No circularity: empirical engineering claims with no derivational reduction

full rationale

The paper introduces TreeFlash as a practical modification to one-shot tree drafters via an added MLP layer and two-stage mechanism to better approximate autoregressive conditioning. All performance claims (12% block efficiency, 9% speedup) are presented as empirical outcomes across tasks and models, with no equations, derivations, fitted parameters renamed as predictions, or self-citations invoked as load-bearing uniqueness theorems. The central assumption about MLP fidelity is stated but not reduced to any input by construction; the method remains an independent engineering proposal whose validity rests on external benchmarks rather than internal definitional equivalence.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Abstract provides no equations, training details, or modeling assumptions; ledger entries cannot be populated from available text.

pith-pipeline@v0.9.1-grok · 5707 in / 1037 out tokens · 21548 ms · 2026-06-28T10:53:16.749033+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

31 extracted references · 5 canonical work pages · 5 internal anchors

  1. [1]

    Evaluating Large Language Models Trained on Code

    Evaluating Large Language Models Trained on Code , author =. arXiv preprint arXiv:2107.03374 , year =

    work page internal anchor Pith review Pith/arXiv arXiv

  2. [2]

    Program Synthesis with Large Language Models

    Program Synthesis with Large Language Models , author =. arXiv preprint arXiv:2108.07732 , year =

    work page internal anchor Pith review Pith/arXiv arXiv

  3. [3]

    Training Verifiers to Solve Math Word Problems

    Training Verifiers to Solve Math Word Problems , author =. arXiv preprint arXiv:2110.14168 , year =

    work page internal anchor Pith review Pith/arXiv arXiv

  4. [4]

    International Conference on Learning Representations , year =

    Let's Verify Step by Step , author =. International Conference on Learning Representations , year =

  5. [5]

    and Zhang, Hao and Gonzalez, Joseph E

    Zheng, Lianmin and Chiang, Wei-Lin and Sheng, Ying and Zhuang, Siyuan and Wu, Zhanghao and Zhuang, Yonghao and Lin, Zi and Li, Zhuohan and Li, Dacheng and Xing, Eric P. and Zhang, Hao and Gonzalez, Joseph E. and Stoica, Ion , booktitle =. Judging. 2023 , url =

    2023

  6. [6]

    Code Alpaca: An Instruction-following

    Chaudhary, Sahil , year =. Code Alpaca: An Instruction-following. GitHub repository , howpublished =

  7. [7]

    2025 , month = aug, url =

    Nathawani, Dhruv and Ding, Shuoyang and Lavrukhin, Vitaly and Gitman, Igor and Majumdar, Somshubra and Bakhturina, Evelina and Ginsburg, Boris and Polak Scowcroft, Jane , version =. 2025 , month = aug, url =

    2025

  8. [8]

    2025 , url =

    Yang, An and Li, Anfeng and Yang, Baosong and Zhang, Beichen and Hui, Binyuan and Zheng, Bo and Yu, Bowen and Gao, Chang and Huang, Chengen and Lv, Chenxu and others , journal =. 2025 , url =

    2025

  9. [9]

    Proceedings of the 40th International Conference on Machine Learning , year =

    Fast Inference from Transformers via Speculative Decoding , author =. Proceedings of the 40th International Conference on Machine Learning , year =

  10. [10]

    2023 , url =

    Sun, Ziteng and Suresh, Ananda Theertha and Ro, Jae Hun and Beirami, Ahmad and Jain, Himanshu and Yu, Felix , booktitle =. 2023 , url =

    2023

  11. [11]

    2024 , publisher =

    Miao, Xupeng and Oliaro, Gabriele and Zhang, Zhihao and Cheng, Xinhao and Wang, Zeyu and Zhang, Zhengxin and Wong, Rae Ying Yee and Zhu, Alan and Yang, Lijie and Shi, Xiaoxiang and Shi, Chunan and Chen, Zhuoming and Arfeen, Daiyaan and Abhyankar, Reyna and Jia, Zhihao , booktitle =. 2024 , publisher =

    2024

  12. [12]

    and Chen, Deming and Dao, Tri , journal =

    Cai, Tianle and Li, Yuhong and Geng, Zhengyang and Peng, Hongwu and Lee, Jason D. and Chen, Deming and Dao, Tri , journal =. 2024 , url =

    2024

  13. [13]

    2024 , url =

    Ankner, Zachary and Parthasarathy, Rishab and Nrusimha, Aniruddha and Rinard, Christopher and Ragan-Kelley, Jonathan and Brandon, William , journal =. 2024 , url =

    2024

  14. [14]

    2024 , publisher =

    Li, Yuhui and Wei, Fangyun and Zhang, Chao and Zhang, Hongyang , booktitle =. 2024 , publisher =

    2024

  15. [15]

    2025 , url =

    Wang, Jikai and Su, Yi and Li, Juntao and Xia, Qingrong and Ye, Zi and Duan, Xinyu and Wang, Zhefeng and Zhang, Min , journal =. 2025 , url =

    2025

  16. [16]

    2025 , url =

    Li, Yuhui and Wei, Fangyun and Zhang, Chao and Zhang, Hongyang , booktitle =. 2025 , url =

    2025

  17. [17]

    2025 , url =

    Liu, Jingyu and Dong, Xin and Ye, Zhifan and Mehta, Rishabh and Fu, Yonggan and Singh, Vartika and Kautz, Jan and Zhang, Ce and Molchanov, Pavlo , journal =. 2025 , url =

    2025

  18. [18]

    2025 , url =

    Li, Guanghao and Fu, Zhihui and Fang, Min and Zhao, Qibin and Tang, Ming and Yuan, Chun and Wang, Jun , journal =. 2025 , url =

    2025

  19. [19]

    and Hartvigsen, Thomas and Fioretto, Ferdinando , journal =

    Sandler, Jameson and Christopher, Jacob K. and Hartvigsen, Thomas and Fioretto, Ferdinando , journal =. 2025 , url =

    2025

  20. [20]

    International Conference on Learning Representations , year =

    Bridging Draft Policy Misalignment: Group Tree Optimization for Speculative Decoding , author =. International Conference on Learning Representations , year =

  21. [21]

    2026 , url =

    Liu, Fuliang and Li, Xue and Zhao, Ketai and Gao, Yinxi and Zhou, Ziyan and Zhang, Zhonghui and Wang, Zhibin and Dou, Wanchun and Zhong, Sheng and Tian, Chen , journal =. 2026 , url =

    2026

  22. [22]

    2026 , url =

    Chen, Jian and Liang, Yesheng and Liu, Zhijian , journal =. 2026 , url =

    2026

  23. [23]

    Accelerating Speculative Decoding with Block Diffusion Draft Trees

    Accelerating Speculative Decoding with Block Diffusion Draft Trees , author =. arXiv preprint arXiv:2604.12989 , year =

    work page internal anchor Pith review Pith/arXiv arXiv

  24. [24]

    International Conference on Learning Representations , volume=

    Distillspec: Improving speculative decoding via knowledge distillation , author=. International Conference on Learning Representations , volume=

  25. [25]

    and Kaiser,

    Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N. and Kaiser,

  26. [26]

    OpenAI GPT-5 System Card

    Openai gpt-5 system card , author=. arXiv preprint arXiv:2601.03267 , year=

    work page internal anchor Pith review Pith/arXiv arXiv

  27. [27]

    Tencent AngelSlim Project Contributors , year=

  28. [28]

    Advances in Neural Information Processing Systems , year =

    FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness , author =. Advances in Neural Information Processing Systems , year =

  29. [29]

    Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song , booktitle =

  30. [30]

    Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles , year =

    Efficient Memory Management for Large Language Model Serving with PagedAttention , author =. Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles , year =

  31. [31]

    2019 , booktitle=

    Decoupled Weight Decay Regularization , author=. 2019 , booktitle=

    2019