RheoSampling: Resolving the One-Hot Dilemma in Stochastic Dynamic-Tree Speculative Decoding
Authors: Qiao Hu, Yepeng Weng, Bo Zhang, Takehisa Yairi
Organizations: National Center for Mathematics and Interdisciplinary Sciences (NCMIS), AMSS, CAS · The University of Tokyo · Lenovo AI Technology Center · SKLMS and AMSS, Chinese Academy of Sciences · School of Mathematical Sciences, University of Chinese Academy of Sciences
Speculative decoding accelerates LLM inference by drafting multiple tokens in parallel, with tree-based methods further improving efficiency through hierarchical structures. Dynamic-tree methods such as EAGLE-3 perform well under greedy decoding via deterministic top-K expansion and global pruning. However, in stochastic decoding (T>0), this mechanism collapses the draft distribution into one-hot probabilities, causing a severe drop in acceptance rate. This creates a dilemma: dynamic-tree methods sacrifice stochastic sampling to preserve context-aware topology, while static-tree methods preserve stochastic sampling with context-agnostic structures. The issue arises because the same probability distribution is used for two conflicting tasks: constructing the tree and verifying tokens. This coupling makes direct injection of randomness challenging due to the resulting stochastic process. We resolve this by decoupling these roles: RheoSampling assigns a token sampled from the draft distribution a proxy probability for tree expansion and pruning alongside its true sampling probability for verification. Specifically, we inject a sampled token among the deterministic top-K slots and treat it with different probabilities during construction and verification, making RheoSampling the first dynamic-tree method with both context-aware top-K construction and stochastic sampling while maintaining losslessness. We establish the lossless guarantee through an equivalence-class analysis that compresses the stochastic tree space into tractable classes. An OT-based verification strategy and a sparse draft mechanism ensure that theoretical gains translate into practical efficiency. Experiments across LLMs and benchmarks demonstrate improvements in acceptance rate and speedup over state-of-the-art dynamic tree methods. This framework may provide a template for analyzing stochastic tree structures.
Figures & tables
Figure 1: Illustration of acceptance rates under different sampling strategies. (a) Sample-based verification: the acceptance rate equals the distribution overlap ∑xmin(p(x),q(x)) . (b) Top- 1 -based verification: the draft deterministically selects xtop1=argmaxxq(x) , yielding an acceptance rate of p(xtop1) , which equals the probability mass the target model assigns to this single point.
Figure 2: Illustrative comparison of multi-draft verification strategies in a simplified two-draft scenario. Starting with the initial top-1 token, the second candidate is either deterministically selected (a) or stochastically sampled (b, c). (a) Top- K : Covers only isolated point masses. (b) RRS-based: Explores the distribution tail but verifies sequentially. (c) OT-based (Ours): Globally reallocates the probability mass (using the normalized draft q~ , resulting in larger overlap with target p ).
Proxy
Sampled Y
Fill tokens
Pool ranking ( R=2 )
Survival tokens
Rheo: 0.5
q~(B)=0.6
C
A(0.5)≥Y(0.5)≥C(0.2)
A,Y←B
q~(C)=0.4
B
A(0.5)≥Y(0.5)≥B(0.3)
A,Y←C
Adhoc: 0.25
q~(B)=0.6
C
A(0.5)≥Y(0.25)≥C(0.2)
A,Y←B
q~(C)=0.4
B
A(0.5)≥B(0.3)≥Y(0.25)
A,B
Table 1: Candidate-pool rankings and survival under two proxy choices.
Model
Sampling
Tasks
Average
Alpaca
Math
Code
MT
QA
Sum
τ
Speedup
V-13B
Top- K
5.59 ±0.09
5.85 ±0.09
6.66 ±0.13
5.69 ±0.06
4.91 ±0.08
5.80 ±0.03
5.75 ±0.03
3.37 ×
Rheo
5.92 ±0.05
5.87 ±0.05
6.72 ±0.06
5.88 ±0.09
5.00 ±0.04
5.94 ±0.07
5.89 ±0.03
3.43 ×
L31-8B
Top- K
5.68 ±0.02
5.50 ±0.03
6.11 ±0.09
4.60 ±0.06
3.89 ±0.04
4.49 ±0.07
5.04 ±0.02
2.84 ×
Rheo
6.03 ±0.12
5.60 ±0.08
6.23 ±0.03
4.84 ±0.14
4.14 ±0.04
4.66 ±0.07
5.25 ±0.03
2.93 ×
DSL-8B
Top- K
4.48 ±0.04
6.70 ±0.05
5.43 ±0.09
4.93 ±0.02
4.08 ±0.05
4.32 ±0.06
4.99 ±0.01
2.89 ×
Table 2: Main results with EAGLE-3 on six different tasks. We report the mean acceptance length (mean ±std over 3 runs) and the end-to-end speedup. V-13B, L31-8B and DSL-8B denote Vicuna-13B-v1.3, Llama-3.1-8B-Instruct and DeepSeek-R1-Distill-Llama-8B, respectively.
Figure 3: Ablation on the rheostat parameter m . Top row: Llama-3.1-8B-Instruct across tasks (GSM8K, HumanEval, Natural Questions). Bottom row: Results on MT-bench across different target models (Llama-3.1-8B-Instruct, DeepSeek-R1-Distill-Llama-8B, Vicuna-13B-v1.3).
Figure 4: Temperature ablation on Llama-3.1-8B-Instruct across MT-bench, GSM8K, and HumanEval. The dashed grey line marks the T=0 (greedy decoding) reference.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Sampling
Verification
Alpaca
GSM8K
HumanEval
MT-bench
Natural Q.
CNN/DM
Average
Top- K
Vanilla
5.68
5.50
6.11
4.60
3.89
4.49
5.04
Rheo
RRSw
5.88
5.47
6.18
4.67
4.09
4.58
5.15
Rheo
6.03
5.60
6.23
4.84
4.14
4.66
5.25
Appendix
Table 3: End-to-end performance comparison across benchmarks. RheoSampling combined with RheoVerification consistently outperforms both the RRSw variant and the vanilla Top- K baseline.
Sampling
Verification
Acceptance
Top- K
Vanilla
80.0%
Rheo
RRSw
83.5%
Rheo
85.4%
Appendix
Table 4: Acceptance rate comparison under controlled settings (MT-bench, draft depth=1, fixed candidate slots: top-1, sampled-1, and top-1 of remaining).
Sampling
Support Size
Coverage of q
τ
Drafting Latency
Rheo
128
91.0%
4.84
13.1 ms
256
91.8%
4.85
13.2 ms
512
92.2%
4.78
13.3 ms
1024
94.1%
4.83
13.3 ms
Full
100.0%
4.89
18.8 ms †
Top-K
N/A
N/A
4.60
12.9 ms
Appendix
Table 5: Effect of sparse support size on draft coverage and average acceptance length ( τ ). The Full row provides a theoretical upper-bound reference using the dense vocabulary distribution. Drafting latency is reported for completeness; the Full configuration uses native PyTorch operators such as concatenation and multinomial, and is not optimized for speed.
Context-aware dynamic trees allocate the speculative decoding budget according to draft path probabilities, adapting their depth and branching to the current context. Under stochastic decoding, however, we find that this structural advantage does not always compensate for the acceptance gains of random sampling paired with advanced verification, and such dynamic trees can fall behind sampled chains in some settings. These trees grow their topology from the candidates themselves, so the tokens submitted for verification are typically the deterministic high-score tokens selected during construction. This coupling is not inherent: once the topology is fixed, its nodes can be repopulated by sampling, allowing dynamic trees to retain their structural advantage while also benefiting from random sampling and advanced verification. Diffusion-based drafters make this practical, as their parallel outputs or lightweight conditional corrections allow candidates to be regenerated cheaply after the complete topology is known. We introduce Tsubame, a two-pass tree speculative decoding framework for diffusion-based drafters. The first pass plans and freezes a context-aware topology using draft path scores; the second replays the fixed topology, sampling the tokens that populate its nodes to form the candidate tree for verification. We prove that Tsubame is lossless under compatible sampling and verification strategies. Experiments across three diffusion-based drafters, six datasets, and multiple candidate budgets show that Tsubame improves acceptance length and throughput over deterministic trees, including settings where it reverses their disadvantage against sampled chains.
Yepeng Weng, Qiao Hu, Takehisa Yairi
The University of Tokyo · National Center for Mathematics and Interdisciplinary Sciences (NCMIS), AMSS, CAS
Multi-path speculative decoding accelerates lossless sampling from a target model by using a cheaper draft model to generate a draft tree of tokens, and then applies a verification algorithm that accepts a subset of these. While prior work has proposed various verification algorithms for i.i.d rollouts, their relative performance under matched settings remains unclear. In this work, we firstly present a systematic evaluation of verification strategies across model families, tasks, and sampling regimes, and find that Traversal Verification dominates consistently, with OT-based methods lagging far behind. Our analysis uncovers that this occurs because OT-based methods achieve high multi-token acceptance near the root of the draft tree, while multi-token gains are most impactful deeper in the draft tree, where draft and target distributions diverge. Based on this insight, we propose delayed tree expansion, which drafts a partial single path, delaying the i.i.d. branching point. We show that delayed tree expansion preserves the target distribution and improves on root-node i.i.d rollouts. Further, we develop a dynamic neural selector that estimates the expected block efficiency of optimal-transport-based verification methods from draft and target features, enabling context-dependent expansion decisions. Our neural selector allows OT-based methods like SpecInfer to outperform Traversal Verification for the first time, achieving 5% higher average throughput across a wide range of models, datasets, and sampling settings.
Rahul Thomas, Teo Kitanovski, Micah Goldblum +1
Ritual · Columbia University · Vanderbilt University
Speculative decoding accelerates language-model inference by letting a cheap drafter propose tokens that the target model verifies in parallel. Recent block drafters make drafting nearly free: a single backbone pass emits an entire block of draft tokens. Draft trees promise a further gain -- several alternative continuations verified in one target forward -- but existing constructions rank candidates by per-position marginals that ignore which parent a candidate extends, so on semi-autoregressive drafters wider trees mostly add mis-ranked nodes; and a tree of fixed size ignores how much speculation each decoding round, and each serving load, can support. We introduce TreeSpark, which reads a parent-conditioned distribution from the drafter's existing Markov head at negligible cost, calibrates it into an edge-acceptance estimate, and lets path survival govern everything else: best-first expansion, per-round stopping, and a load-adaptive serving policy. Sampling siblings without replacement, with matching residuals in recursive rejection, keeps decoding lossless at any temperature. Adaptive trees improve on matched fixed budgets at every temperature; against a tuned chain on the same drafter, TreeSpark accepts 15-25% more draft tokens per round and decodes 8-14% faster in single-request wall-clock, and under rising load it gracefully shrinks the tree back to the chain. Code and artifacts: https://github.com/PopSoda2002/TreeSpark
Huapeng Zhou, Huayu Wang, Xinyu Wang
Boson AI · University of Washington · McGill University