Speculative Beam Search

The core mechanism uses a small, fast draft model to generate multiple candidate continuations (a speculative candidate sequence) for each active beam in parallel. This replaces the standard autoregressive expansion of each beam, which is sequential and costly. The draft model proposes γ (gamma) tokens ahead for every beam, creating a tree of possibilities that is then efficiently verified in a single batch by the target model.

• Key Benefit: Replaces B * γ sequential target model calls with 1 parallel verification pass, where B is the beam width.
• Example: With a beam width of 4 and a speculative factor of 5, the draft model generates 20 candidate tokens in parallel, which are verified in one batched forward pass.

The target model (the large, accurate LLM) performs a single, batched verification forward pass to score all speculative tokens proposed across all beams. This pass uses a modified tree attention mechanism to compute logits for each candidate position efficiently. The model compares its own predicted token distribution against the draft's proposals.

• Verification Cost: The computational cost of this single batched pass must be less than the cost of the target model generating B * γ tokens autoregressively for a net speedup.
• Process: For each position, if the target model's predicted probability for the draft token is higher than a random threshold, the token is accepted; otherwise, it is rejected, and the target model's own token is sampled.

After verification, sequences are scored and the beam is pruned. The score for a beam is a function of:

• The cumulative log probability of all accepted tokens.
• The target model's probability for any corrected tokens (where the draft was rejected).
• A length normalization factor.

Beam pruning then selects the top-B highest-scoring sequences from the expanded tree of possibilities. This integrates speculative efficiency directly into the beam search's exploration-vs-exploitation trade-off, maintaining output quality while reducing latency.

When a draft token is rejected, a rollback mechanism is triggered for that specific beam path. Generation reverts to the last accepted position. The target model's own sampled token is used, and generation continues autoregressively from that point for the remainder of the step. This ensures the output distribution remains identical to that of the standard target model, preserving accuracy.

• Early Stopping: Verification for a beam can halt upon the first rejection to save computation, a common optimization.
• State Management: The speculative KV cache for accepted tokens is retained; the cache for rejected speculative branches is discarded.

The technique navigates a three-way optimization:

• Latency/Throughput: Speedup is maximized with a high acceptance rate from the draft model and optimal choice of beam width (B) and speculative factor (γ).
• Output Quality: Controlled by the beam width and the accuracy of the target model's verification. The latency-accuracy tradeoff is minimal as the final output distribution is governed by the target model.
• Hardware Efficiency: Performance depends on hardware-aware speculation. The optimal γ is often limited by GPU memory bandwidth for the large batched verification pass.

Speculative beam search is a direct evolution of core speculative decoding concepts:

• vs. Standard Speculative Decoding: Operates on multiple beams (B > 1) instead of a single sequence, improving output diversity and quality for tasks like translation or summarization.
• vs. Lookahead Decoding: Uses a separate, trained draft model rather than n-gram tables or the model's own heads, typically yielding higher acceptance rates.
• vs. Self-Speculative Decoding (e.g., Medusa): Employs a distinct draft model, avoiding the parameter overhead of adding Medusa heads to the target model but requiring a separate model load. Its primary advantage is applying speculative speedups to beam search, a previously sequential bottleneck.

Beam Search is a heuristic search algorithm that expands the most promising nodes in a limited set (the beam width) at each step of a sequence generation task. Unlike greedy decoding, which commits to the single highest-probability token at each step, beam search maintains multiple candidate sequences (beams) in parallel.

• Core Mechanism: At each decoding step, it generates all possible next-token extensions for all active beams, scores the resulting sequences, and keeps only the top-k highest-scoring sequences.
• Trade-off: It provides a better approximation of the globally optimal sequence than greedy search but increases memory and compute cost proportional to the beam width.
• Use in SBS: Speculative Beam Search uses beam search within the draft model to generate multiple high-likelihood candidate continuations for each beam before verification.

A Draft Model is a smaller, faster language model used in speculative decoding techniques to propose candidate token sequences. Its primary function is to approximate the behavior of a larger, more accurate target model at a fraction of the computational cost.

• Characteristics: Typically has significantly fewer parameters, enabling faster autoregressive generation. It is often a distilled version of the target model or a model from the same family (e.g., using a 7B parameter model to draft for a 70B parameter model).
• Role in SBS: In Speculative Beam Search, the draft model performs beam search to generate k * gamma candidate tokens (where k is beam width, gamma is lookahead), creating a tree of potential continuations for each active beam in the target model.

Tree Attention is a modified attention mechanism that enables a transformer model to process a tree of token sequences in a single, batched forward pass. It is a critical enabling technology for efficient verification in speculative decoding with multiple candidates.

• How it Works: It manages attention masks and positional encodings to allow tokens in different candidate branches to attend to their correct, shared prefix history while preventing cross-talk between separate branches.
• Application: In Speculative Beam Search, Tree Attention allows the target model to verify all tokens across all candidate beams and lookahead positions simultaneously. This parallel verification is what makes the speculative approach computationally viable versus sequential autoregressive scoring.

The Verification Forward Pass is the single, batched inference step where the target model scores all tokens in the speculative candidate sequences proposed by the draft model.

• Process: The target model takes the batched tree of candidate tokens (from all beams and lookahead positions) and computes logits for each position. It then compares its own predicted token distribution against the draft's proposed tokens.
• Acceptance Rule: A draft token is accepted if it matches the target model's most likely token (greedy) or, in a probabilistic variant, if it is sampled from the adjusted target distribution.
• Cost Analysis: The computational cost of this single large batched pass must be less than the cost of the target model generating the same number of tokens autoregressively for a net speedup.

A Rollback Mechanism is the procedure invoked when the target model rejects a token proposed by the draft model during speculative verification. It ensures the output sequence remains identical to what the target model would have generated autoregressively.

• Trigger: Occurs at the first position in a candidate sequence where the draft token is rejected.
• Action: Generation reverts to the last unanimously accepted token. The target model then generates the correct token autoregressively from that point, and the speculative process restarts.
• In SBS: In Speculative Beam Search, rejection and rollback can happen independently for each beam. A beam where a token is rejected early may fall behind other beams that accepted more draft tokens, dynamically affecting the search paths.

The Latency-Accuracy Trade-off in speculative decoding refers to the balance between inference speedup and strict adherence to the target model's original output distribution. Speculative methods introduce a risk of distributional shift.

• Speedup Source: Latency reduction comes from replacing multiple slow target model steps with fewer batched verification steps, amortizing the cost of the large model.
• Accuracy Risk: Using an imperfect draft model introduces bias. The acceptance/rejection mechanism mitigates this but does not eliminate it entirely, as the draft influences which tokens are presented for verification.
• SBS Specifics: In Speculative Beam Search, parameters like beam width (k) and speculative lookahead (gamma) directly control this trade-off. Larger k and gamma increase parallelism and potential speedup but also amplify the draft model's influence on the search space exploration.