lonerapier.me

Speculative Decoding

4 min read

  • Eagle
    • Eagle series pushes a step further in complexity spectrum and operates at the feature level. Rather than predicting the next token, EAGLE hypothesizes that feature sequences exhibit more regularity and attaches a lightweight autoregressive drafter predicting the next feature. This skips the complexity that arrives from training separate prediction heads as in Medusa. Draft layer (autoregressive head) consisting takes the embedding of the predicted token (from the base model at the start or the last draft candidate), along with the features of the previous generated token. The autoregressive head predicts the next feature. The LM head as per the feature, predicts the distribution. The distribution is then used to sample the next token. The predicted feature and sampled token are concatenated, and added as input for the next autoregressive draft phase.
  • Eagle-2
    • Observes that acceptance of draft candidates, besides the position also depends on the context. Static draft trees generated for verification is same for all context ( position generates 2 tokens, generates 3). Generating dynamic draft trees by obtaining confidence scores from the draft model greatly improves acceptance rates. Using the confidence score from draft model is fine because draft model is well-calibrated, and approximates acceptance rates from the target distribution. This enables EAGLE head to chain obvious outputs into longer predicted sequences, dynamically adjusting the length as per the context.
  • Eagle-3
    • Drops feature prediction in favor of direct token prediction, simplifying the EAGLE’s loss function in the process. Moreover, instead of using just the final hidden layer representation which are extremely focused on next-token and forces the drafter heads to re-derive the context, EAGLE fuses the feature representations from the target model’s decoder layers and directly predicts the next feature using a lightweight autoregressive prediction head. Simulating multi-step generation by feeding draft models its own output at training time minimises the training/inference gap which the paper refers to as training-time test. Uses dynamic draft trees from EAGLE-2 to propose multiple chained continuations, and efficient parallel tree attention from EAGLE to effectively prune invalid tree branches in one parallel pass.
  • EAGLE refines the speculative decoding procedure at every level (context, prediction, verification), but is still autoregressive, and suffers from diminishing gains as the draft length grows.
  • MTP
  • DFlash: Using a block-dLLM to denoise a block in arbitrary manner while generating blocks autoregressively.
  • DSpark

Next things:

  • Hands on tutorial for speculative decoding using rejection sampling approach
  • Switch to DFlash style diffusion models
  • Alignment: How draft models are aligned with target?
  • Distillation: How much can draft models be compressed?

References