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

Speculative decoding accelerates LLM inference by drafting multiple tokens in parallel, with tree-based methods further improving efficiency through structured hierarchies. Dynamic-tree methods such as EAGLE-3 achieve excellent performance 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 exposes an apparent 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 the tokens. This coupling makes direct injection of randomness extremely challenging, as we are faced with a complex stochastic process. We resolve this by decoupling these two 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 in the construction and verification process, 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 a novel 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 diverse LLMs and benchmarks demonstrate consistent improvements in acceptance rate and speedup over state-of-the-art dynamic tree methods. This framework may provide a template for analyzing other complex stochastic tree structures.