RheoSampling speeds up large language model predictions with better random sampling
RheoSampling: Resolving the One-Hot Dilemma in Stochastic Dynamic-Tree Speculative Decoding
Computation and LanguageMachine Learning
Summary
When computers generate text, they often try many word guesses in parallel to speed up prediction. There are smart ways to organize these guesses in a tree structure, but this gets tricky when the computer needs to pick words with randomness to sound more natural. The authors show that trying to use one probability for both building the guess tree and picking words caused problems. Their new method, RheoSampling, uses separate probabilities for tree building and word verification, allowing the model to keep randomness and context awareness without losing accuracy. This improves the speed and quality of generating text with large language models.
What this means in practice
- •For machine learning engineers: Improve inference speed and flexibility of large language models in generating varied text outputs.
- •For conversational ai developers: Enhance chatbot response diversity and efficiency by enabling fast stochastic decoding with context-aware token selection.
Authors
Qiao Hu, Yepeng Weng, Bo Zhang, Takehisa Yairi
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.