RheoSampling: Resolving the One-Hot Dilemma in Stochastic Dynamic-Tree Speculative Decoding

Qiao Hu; Yepeng Weng; Bo Zhang; Takehisa Yairi

cs.CLSep 18, 2026

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 ·

Abstract

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 $\sum_{x}\min(p(x),q(x))$ . (b) Top- $1$ -based verification: the draft deterministically selects $x_{top1}=\arg\max_{x}q(x)$ , yielding an acceptance rate of $p(x_{top1})$ , 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 $\tilde{q}$ , resulting in larger overlap with target $p$ ).

Proxy	Sampled $Y$	Fill tokens	Pool ranking ( $R=2$ )	Survival tokens
Rheo: $0.5$	$\tilde{q}(B)=0.6$	$C$	$\boxed{A(0.5)\geq Y(0.5)}\geq C(0.2)$	$A,Y\leftarrow B$
Rheo: $0.5$	$\tilde{q}(C)=0.4$	$B$	$\boxed{A(0.5)\geq Y(0.5)}\geq B(0.3)$	$A,Y\leftarrow C$
Adhoc: $0.25$	$\tilde{q}(B)=0.6$	$C$	$\boxed{A(0.5)\geq Y(0.25)}\geq C(0.2)$	$A,Y\leftarrow B$
Adhoc: $0.25$	$\tilde{q}(C)=0.4$	$B$	$\boxed{A(0.5)\geq B(0.3)}\geq Y(0.25)$	$A,B$

Table 1: Candidate-pool rankings and survival under two proxy choices.

Model	Sampling	Tasks						Average
Model	Sampling	Alpaca	Math	Code	MT	QA	Sum	$\tau$	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 $\times$
V-13B	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 $\times$
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 $\times$
L31-8B	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 $\times$
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 $\times$

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	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	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	$\boldsymbol{\tau}$	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 (

\tau

). 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.

Speculative Decoding

Speculative decoding accelerates generation by drafting multiple tokens and verifying them in parallel with a larger target model.

See all 311 papers

Tsubame: Tree Replay for Diffusion-Based Speculative Decoding
Sep 27, 2026Yepeng Weng, Qiao Hu, Takehisa YairiSpeculative Decoding for Diffusion Language Models Speculative Decoding
The University of Tokyo · National Center for Mathematics and Interdisciplinary Sciences (NCMIS), AMSS, CAS
Dynamic Delayed Tree Expansion For Improved Multi-Path Speculative Decoding
Feb 19, 2026Rahul Thomas, Teo Kitanovski, Micah Goldblum +1Speculative Decoding LLM Inference Acceleration
Ritual · Columbia University · Vanderbilt University
TreeSpark: Calibrated, Load-Adaptive Draft Trees for Semi-Autoregressive Speculative Decoding
Aug 12, 2026Huapeng Zhou, Huayu Wang, Xinyu WangSpeculative Decoding LLM Inference Acceleration
Boson AI · University of Washington · McGill University

Sep 27, 2026J/K move · Enter open · S save

Tsubame: Tree Replay for Diffusion-Based Speculative Decoding

Yepeng Weng, Qiao Hu, Takehisa Yairi

The University of Tokyo · National Center for Mathematics and Interdisciplinary Sciences (NCMIS), AMSS, CAS

Speculative Decoding for Diffusion Language Models Speculative Decoding Diffusion Language Model Inference Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

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.

Abstract

Speculative Decoding

Explore similar work

Tsubame: Tree Replay for Diffusion-Based Speculative Decoding

Dynamic Delayed Tree Expansion For Improved Multi-Path Speculative Decoding

TreeSpark: Calibrated, Load-Adaptive Draft Trees for Semi-Autoregressive Speculative Decoding