Beam search, speculative decoding & Medusa

1

Concept

Autoregressive generation exposes a serial dependency: token $t+1$ cannot be finalized until token $t$ is known. Several methods spend extra parallel work around this chain, but they optimize different things. Beam search explores sequences to improve a search objective. Speculative decoding and Medusa propose future tokens so the expensive model may validate several positions in one pass.

Greedy decoding keeps one partial sequence. Beam search keeps $b$ candidates. At each step it expands each beam with possible next tokens, adds log-probabilities to cumulative scores, and retains the best candidates according to a scoring rule. Length normalization and stopping rules are essential because raw summed log-probability favors shorter sequences. Beams can share KV-cache prefixes and branch through copy-on-write metadata.

Beam search is useful when a concentrated output and sequence-level likelihood matter, such as some constrained transduction tasks. For open-ended dialogue, several beams often collapse into near-duplicates or favor bland, high-probability text. Increasing beam width raises model work and cache use. It is not a general inference acceleration technique.

Speculative decoding starts with a smaller or cheaper draft model. The draft autoregressively proposes a block of future tokens. The larger target model evaluates those proposed positions together, exploiting parallelism similar to a short prefill. An acceptance procedure compares draft and target probabilities. Accepted tokens advance the sequence; after the first rejection, dependent draft suffix tokens are discarded and a corrected token is sampled.

With the proper acceptance and correction rule, speculative sampling preserves the target model’s output distribution. This is stronger than merely checking whether the target’s argmax equals the draft token. Exactness refers to the mathematical sampler under matching numerical and implementation assumptions, not bit-identical output across different kernels.

Speed depends on acceptance rate and relative cost. A draft that is too weak proposes many rejected tokens. A draft that is too expensive saves little. Longer proposal blocks offer more potential progress but waste more work after early rejection. The best configuration depends on prompts, decoding settings, batch size, hardware, and load.

Medusa avoids a separate full draft model by attaching multiple lightweight decoding heads to the base model. Heads predict tokens at several future offsets, creating a tree of candidate continuations. A verification pass checks the tree with the base model and accepts a valid path. Training or adapting these heads must match the target checkpoint; changing the base can invalidate their proposal quality.

Tree attention and cache bookkeeping are the quiet complexity. Candidate branches share a prefix but differ later. Masks must prevent one branch from reading another. Accepted states must be committed in order; rejected branches must release memory. A bug here can produce plausible but incorrect tokens, so distributional tests against baseline decoding matter.

Constrained decoding can combine with these methods, but proposal and verification must obey the same grammar state. Tool calls and stop sequences also change acceptance boundaries. Streaming may arrive in bursts when several tokens are accepted, so clients must not assume one network event per model step.

The durable distinction is purpose. Beam search pays for multiple hypotheses to choose a preferred sequence. Speculation pays a cheap proposer to expose parallel work while the target remains authoritative. Medusa moves proposal heads into the target architecture. All three branch into possible futures; only their objective, acceptance contract, and cost decide whether that branching helps.