Draft Model

The draft model's fundamental role is to perform fast, autoregressive token generation to create a candidate sequence (typically 3-10 tokens, defined by the speculative factor, γ). This sequence is a 'best guess' of what the target model would generate next. The draft operates with minimal latency, trading off some accuracy for raw speed, as its outputs are provisional and subject to verification.

Draft models are architecturally optimized for minimal latency and memory footprint, not maximum accuracy. Common characteristics include:

• Reduced Parameters: Often 10x-100x smaller than the target model (e.g., a 125M parameter draft for a 7B parameter target).
• Shallow Layers: Fewer transformer blocks to decrease sequential computation.
• Optimized for Single-Token Latency: The primary metric is time-to-first-token and the time to generate the γ-token draft sequence, which must be less than the target's time to generate those same tokens autoregressively.

To maximize the acceptance rate, the draft model must be distributionally aligned with the target model. Common training methodologies include:

• Distillation: Training the draft on the logits or output distributions of the target model on a broad corpus.
• Domain-Specific Tuning: Fine-tuning the draft on data representative of the target's deployment context to improve prediction relevance.
• Architectural Homogeneity: Using a smaller version of the same model family (e.g., Llama 3 8B as draft for Llama 3 70B) can improve alignment due to shared tokenizers and training data.

A draft model's efficacy is measured by its impact on the overall inference system:

• Acceptance Rate: The percentage of proposed tokens accepted by the target model. A higher rate directly increases the speedup factor. Rates of 70-85% are common for well-aligned pairs.
• Speedup Factor: The ratio T_autoregressive / T_speculative. Net speedup occurs when: (Time to Draft γ tokens) + (Time to Verify γ tokens) < (Time for Target to generate γ tokens autoregressively).
• Verification Cost: The computational overhead of the target model's parallel scoring pass. The draft must be fast enough that this cost doesn't erase the gains.

The draft model is not used in isolation; it is a component in a tightly coupled inference loop:

1. Drafting Phase: The draft generates γ candidate tokens autoregressively, conditioned on the current context.
2. Verification Phase: The target model performs a single, batched forward pass to score all γ tokens in parallel.
3. Accept/Reject & Rollback: Tokens are accepted where the target's probability matches the draft's prediction. On first rejection, a rollback mechanism reverts to that position, and the target model generates the correct token.
4. Cache Management: The speculative KV cache from drafting may be reused to optimize the verification pass.

Beyond a standalone small LM, the 'draft' function can be implemented through other mechanisms:

• Self-Speculative Decoding: The target model drafts for itself using auxiliary structures like Medusa heads, eliminating a separate model.
• N-Gram Drafting: Uses a static table of frequent token sequences from training data for ultra-fast, memory-based proposals.
• Dynamic Draft Selection: A system uses multiple draft models or strategies, selecting the one with highest predicted confidence for the given context.
• Lookahead Decoding: Leverages the target model's own intermediate representations to draft future tokens.

The draft model's architecture is chosen for minimal latency, not maximum accuracy. Common choices include:

• Distilled variants of the target model (e.g., a 1.3B parameter version of a 70B model).
• Specialized small models like Phi-2 or TinyLlama, which offer strong performance per parameter.
• N-gram or lookup-based models for extreme speed but lower acceptance rates.

Draft model distillation is a critical optimization, where the small model is explicitly trained to mimic the next-token probability distributions of the larger target model, directly maximizing the acceptance rate.

The acceptance rate is the primary metric for draft model efficacy. A low rate negates the speedup from parallel verification. Optimization strategies include:

• Temperature tuning: Adjusting the draft model's sampling temperature to better align with the target model's distribution.
• Confidence thresholding: Only proposing a draft token if the draft model's probability exceeds a set threshold (e.g., 0.8), otherwise letting the target model generate.
• Context-aware drafting: Using the recent generation history to predict when the draft model is likely to be accurate (e.g., for common phrases) versus inaccurate (e.g., for novel reasoning).

The speculative factor (γ) defines how many tokens the draft model generates ahead of the target model. This is a crucial hyperparameter with a direct trade-off:

• Small γ (e.g., 3-5): Lower risk. If most tokens are rejected, the rollback cost is minimal. Best for low-confidence domains.
• Large γ (e.g., 8-10): Higher potential speedup if acceptance is high, but a single early rejection wastes the computation on all subsequent drafted tokens.

Optimal γ is found empirically and depends on the draft/target model pair, hardware, and batch size. It must satisfy the core inequality: Cost(verify γ tokens) < γ * Cost(generate 1 token autoregressively).

The draft model must reside in fast memory alongside the target model to avoid I/O bottlenecks. Key considerations:

• GPU Memory Footprint: The draft model, its KV cache, and the target model's expanded cache for verification must fit simultaneously. This often limits the practical size of the draft model.
• Kernel Fusion: Draft model inference must use optimized kernels to ensure its forward pass is truly negligible compared to the target model's.
• Batch Processing: In a batched serving environment, the draft model must generate candidate sequences for all requests in the batch efficiently. Dynamic batching strategies must account for the two-stage draft-verify process.

Efficient integration minimizes overhead. This involves:

• Shared Tokenizers & Embeddings: Using aligned vocabularies to avoid costly re-tokenization or projection between draft and target token spaces.
• KV Cache Management: The speculative KV cache generated during drafting must be efficiently integrated into the target model's attention mechanism for the parallel verification pass. This requires careful memory layout planning.
• Early Stopping & Rollback: The system must implement a low-overhead early stopping mechanism to halt verification upon the first token rejection and a fast rollback to revert to the last accepted token before resuming generation (often from the target model).

Advanced systems may employ multiple draft strategies and select between them dynamically to maximize performance.

• Model Router: A lightweight classifier decides per-token or per-sequence whether to use a fast n-gram draft, a small neural draft, or no draft at all.
• Confidence-Based Fallback: If the primary draft model's confidence is low, the system can bypass drafting for that step, invoking the target model directly to preserve output quality.
• Cost-Aware Drafting: The system can adjust the drafting strategy based on current server load and latency SLAs, trading some speedup for consistency under load.

The target model is the primary, larger, and more accurate language model in speculative decoding. Its role is to verify the sequence of tokens proposed by the draft model in a single, parallel forward pass. The target model's computational cost is the primary bottleneck, which speculative decoding aims to circumvent.

• Function: Acts as the verifier and final arbiter of token correctness.
• Trade-off: Typically has higher parameter count and better accuracy than the draft model, but is slower per autoregressive step.

Token verification is the critical process where the target model evaluates the draft model's proposed candidate sequence. Instead of generating tokens one-by-one, the target model performs a single, batched forward pass over the entire candidate sequence, comparing its own predicted probability distributions against the draft's proposals.

• Mechanism: Uses a modified attention mask to process the candidate sequence in parallel.
• Output: For each position, the target model either accepts the draft token or rejects it and samples from its own distribution.

The acceptance rate is the percentage of tokens proposed by the draft model that are accepted by the target model during verification. This is the single most important metric determining the efficacy of speculative decoding.

• Impact: A high acceptance rate means most draft tokens are correct, maximizing the speedup. A low rate leads to frequent rollbacks, negating benefits.
• Typical Range: Effective systems often require acceptance rates above 70-80% to achieve meaningful speedups (e.g., 2-3x).

Lookahead decoding is a variant of speculative decoding that eliminates the need for a separate, trained draft model. Instead, it uses the target model's own internal mechanisms or external statistical tables to propose candidate tokens.

• N-Gram Drafting: Uses a static table of frequent token sequences (n-grams) from the training corpus.
• Self-Drafting: Employs auxiliary prediction heads (like Medusa heads) attached to the target model to propose multiple future tokens in the same forward pass.

A small-big model pair refers to the specific combination of models chosen for speculative decoding. The pairing is a deliberate engineering trade-off between draft speed and target accuracy.

• Selection Criteria: The draft model must be significantly faster (10-100x) per token to offset verification cost. Common pairs include a 7B parameter model drafting for a 70B model, or a distilled version of the target.
• Architecture Alignment: Ideally, the draft and target share the same tokenizer and have similar output distributions to maximize acceptance rate.

The verification forward pass is the single, batched inference step through the target model that scores all tokens in the speculative candidate sequence. Its computational cost must be less than the cost of generating the same number of tokens autoregressively for a net speedup.

• Key Optimization: Uses a speculative KV cache, where the draft model's cached key-value pairs for the candidate sequence are reused to avoid recomputation.
• Parallelism: The entire candidate sequence (e.g., 3-5 tokens) is processed in parallel, unlike sequential autoregressive decoding.