RheoSampling keeps temperature sampling inside a dynamic draft tree — Proxy tree probabilities vs true verification probabilities — What does it mean?

The news. On September 18, 2026, a paper titled RheoSampling: Resolving the One-Hot Dilemma in Stochastic Dynamic-Tree Speculative Decoding was posted to arXiv. It does not propose a faster drafter — it removes a choice nobody should have to make: until now you could have a draft tree that reshapes itself to the context, or draft candidates genuinely sampled at temperature, but not both. A dynamic draft tree ranks and prunes its nodes using the draft model's own probabilities, and above temperature zero that ranking stops predicting what the target model will accept. The reported end-to-end gains over a top-K baseline are modest, on the order of 0.14 to 0.22 extra tokens per draft-verify cycle, but they arrive with a proof that the sampled output distribution is unchanged. Read the paper →

Picture a casting director with three callback slots and a long list of auditions. Slot one goes to the highest score, no argument. Slot three goes to the next name down. Slot two is a wildcard: a name drawn at random from everyone except the top-scoring one, so an unlikely candidate occasionally gets a hearing.

Now the trap. Suppose the wildcard has to hold its place on the board the same way everyone else does, by its own audition score. A wildcard drawn from near the bottom scores badly, gets bumped off the board, and never reaches the judges. Only the lucky wildcards that happened to score well survive, so the wildcards the judges actually see are no longer a fair draw from the list — they are a draw conditioned on having scored well. Asking one number to both seat the candidate and judge them has quietly rigged the lottery.

RheoSampling's answer is to stop asking one number to do two jobs. The wildcard's place on the board is decided before the draw — slot two, whoever comes out — so surviving the board has nothing to do with who was drawn. The judges then score that person on their real audition.

Underneath the metaphor, the callback board is one level of a draft tree. Speculative decoding runs a small drafter ahead of the big model and checks its guesses in one batched forward pass, and a tree of guesses beats a single chain because a wrong turn on one branch still leaves a branch the verifier can accept. Dynamic tree methods — EAGLE-2 and EAGLE-3 are the ones this paper builds on — rebuild that tree at every step: expand the top-K children of each parent, score each root-to-leaf path by multiplying the drafter's probabilities along it, then prune down to a global node budget.

Two probabilities matter from here on: write q(x) for the probability the small drafter assigns a token, and p(x) for the probability the big target model assigns it. The tree is ranked entirely on q. That ranking is the whole reason a dynamic tree beats a fixed one, and it is exactly what breaks when the temperature goes up. Whether the target accepts a token depends on the ratio p(x)/q(x), not on the size of q(x) by itself — a token the drafter rates as unlikely can still be accepted if the target likes it, and a large q(x) guarantees nothing. At temperature zero the drafter's favourite and the target's favourite agree closely enough that nobody notices the difference. Above zero they come apart.

The paper names the consequence the one-hot collapse. Because every node enters the tree through a deterministic top-K pick, the proposal distribution the tree represents is a set of point masses, and the long tail of q is discarded during expansion. The tail is precisely where a sampled decode spends its time, so the acceptance rate drops.

The obvious repair — sample the nodes instead of taking the top-K — fails for a subtler reason, and this is the part worth slowing down for. When a sampled token is ranked by its own probability, whether it survives pruning depends on which token it turned out to be. Conditioning on survival then bends that token's distribution away from the one it was drawn from. Verification is still doing its accept-or-reject arithmetic against the original distribution, so it is now comparing against the wrong thing, and the guarantee that speculative decoding emits exactly the same text distribution as plain sampling — its losslessness — is gone. A serving change that silently alters what the model writes is not a speedup; it is a different model.

The paper's fix is to give the probe two identities. Of the K slots at each expansion step, the first m go to the highest-probability lead tokens, one slot goes to a token sampled from the residual tail, and the remainder are filled with the next-highest-ranked tokens. The sampled probe then carries two numbers: a proxy probability of min{q(x_m), z}, where x_m is the m-th highest-probability draft token and z is the mass left over after the top-m, used only for expansion, reranking and pruning; and its true sampling probability under that residual tail distribution, used only at verification. (At m = 0 there is no top-m to sit behind, so the paper anchors the probe in the first slot instead, with a proxy just above the largest draft probability.) The proxy is constructed so the probe outranks every fill token, which makes its survival independent of which token it turned out to be — the exact property the naive repair destroyed. That independence is what the paper's losslessness proof rests on: it groups the enormous space of trees the randomness could produce into equivalence classes whose members behave identically under verification, which makes the space small enough to reason about.

Worked example — where the gain goes, and where it leaks. Hold the setup fixed: Llama-3.1-8B-Instruct as the target, the released EAGLE-3 draft head with no extra fine-tuning, a 60-node tree at draft depth 8, temperature T = 1.0, a single NVIDIA A6000, averaged over six 80-question benchmarks and three seeds.

The paper reports an acceptance length — tokens kept per draft-verify cycle — of 5.04 for the top-K baseline and 5.25 for RheoSampling. That is 0.21 / 5.04 = 4.2% more tokens per cycle. Turned into work the server does, emitting 1,000 tokens needs 1000 / 5.04 = 198 verification cycles before and 1000 / 5.25 = 190 after: about eight fewer full forward passes through the 8B target per 1,000 tokens.

End-to-end speedup moves less, 2.84× to 2.93×, which is 2.93 / 2.84 - 1 = 3.2%. The paper attributes the gap between 4.2% and 3.2% to the probe's own overhead — drawing it and verifying it separately is work the pure top-K baseline never does. So roughly three-quarters of the acceptance gain survives into wall-clock, and the other quarter pays for the mechanism. The two other targets move the same way: Vicuna-13B goes 5.75 → 5.89 and 3.37× → 3.43×, DeepSeek-R1-Distill-Llama-8B goes 4.99 → 5.21 and 2.89× → 2.99×.

The rheostat is the knob you actually tune. m trades the probe's survival against the quality of the tree. A small m hands the probe a larger proxy and a higher chance of surviving pruning, but strips out the deterministic backbone the tree was ranking on. A large m keeps the topology clean and shrinks the probe's proxy, which the paper says increases the risk that the probe is pruned away before it can do anything. The paper reports m = 1 as best or near-best in most settings, with m ≥ 3 yielding only marginal gains.

Two other costs are worth naming before anyone reaches for this. The drafter's logits are truncated to the top 128 before the softmax so the sampling step never walks the full vocabulary — lossless only because verification uses that same truncated distribution. And as temperature falls toward zero the sampling step converges on deterministic selection anyway — at m = 0 the tree's shape collapses back to plain top-K too, while at m = 1 the paper reports the topology stays distinct — and the paper notes that the rival verification strategies become nearly equivalent in that limit. So this is a fix for serving stacks that sample, not for every stack.

Goes deeper in: LLM Serving → Speculative Decoding → The Verification Algorithm

Related explainers

• JetSpec — parallel tree drafting — the case for a tree over a chain in the first place. Read it first if the idea of a branching draft is new; this one assumes the tree and asks what temperature does to it.
• DominoTree — conditional draft-tree scoring — the other way to improve a draft tree: change how paths are scored rather than what is allowed into the candidate pool.
• LoopSpec — pipelined self-speculative decoding — a speedup that comes from overlapping draft and verify rather than from the tree, and a useful contrast for where the remaining wall-clock actually sits.