DIVERSED: Relaxed Speculative Decoding via Dynamic Ensemble Verification

1. Problem Statement

Large Language Models (LLMs) face significant inference latency due to their autoregressive nature, where tokens are generated one by one. Speculative Decoding (SD) is a standard acceleration technique that uses a small "draft" model to propose multiple tokens, which a larger "target" model verifies in parallel.

However, standard SD suffers from a verification bottleneck:

• Rigid Verification: The target model strictly enforces that the accepted token distribution must exactly match the target model's distribution.
• Low Acceptance Rates: This strictness leads to the rejection of many plausible tokens generated by the draft model, limiting the speedup.
• Static Relaxation Limitations: Previous "lossy" approaches attempt to relax this rule using a static ensemble (a fixed weighted average of the draft and target distributions). While this improves acceptance rates, a single fixed weight cannot adapt to varying contexts. Some tokens are critical for correctness (e.g., in math), while others are less influential. A static rule either rejects too many good tokens or accepts too many bad ones, degrading generation quality.

2. Methodology: DIVERSED

The authors propose DIVERSED (DynamIc VErification RElaxed SpEculative Decoding), a framework that replaces the rigid verification step with a dynamic ensemble verifier.

Core Mechanism

Instead of a fixed weight, DIVERSED learns a context-dependent and token-dependent weight ($w_t$) to blend the draft model distribution ($q_t$) and the target model distribution ($p_t$).

1. Dynamic Ensemble Distribution:
   At each time step $t$, the verification distribution $\nu_t$ is defined as:
   $$\nu_t(x) = w_t p_t(x) + (1 - w_t) q_t(x)$$
   where $w_t$ is a scalar between 0 and 1.

2. Ensemble Head:
   The weight $w_t$ is predicted by a lightweight neural network (an "ensemble head") that takes as input the hidden representations from both the draft model ($h^q_t$) and the target model ($h^p_t$). This allows the system to assess the reliability of the draft token based on the specific context and the agreement between the two models.

3. Training via Reinforcement Learning:
   Since token-level supervision is often unavailable, DIVERSED is trained using Reinforcement Learning (specifically REINFORCE++):

   • Objective: Maximize a regularized objective that balances Task Reward ($R$) and Acceptance Rate.
   • Loss Function:
     $$\mathbb{E} [ R(x_{1:T}) + \lambda \sum_{t=0}^{T} (1 - TV(q_t, \nu_t)) ]$$
     • $R(x_{1:T})$: Task-specific reward (e.g., accuracy for math, ROUGE for summarization).
     • $TV(q_t, \nu_t)$: Total Variation distance, acting as a proxy for acceptance probability.
     • $\lambda$: A hyperparameter controlling the trade-off between quality and speed.
   • This encourages the verifier to be aggressive (high $w_t$ for target, low acceptance) on quality-critical tokens and lenient (low $w_t$, high acceptance) on less critical tokens.

4. Inference Pipeline:
   During inference, the draft model proposes $N$ tokens. The target model verifies them in parallel. The ensemble head computes dynamic weights for each proposed token, calculating a relaxed acceptance probability:
   $$b_{\nu_t}(\tilde{x}_t) = \min\left(1, \frac{\nu_t(\tilde{x}_t)}{q_t(\tilde{x}_t)}\right)$$
   If rejected, a fallback distribution is sampled.

3. Key Contributions

Theoretical Contributions

• Exact Expected Acceptance Length: The authors derive a precise mathematical expression for the expected number of accepted tokens ($E[T_N]$) without the independent and identically distributed (i.i.d.) assumption used in prior work. They show that acceptance probability depends on the position in the sequence and the sequential dependency of accepted tokens (following a "min-kernel" distribution).
• Pareto Frontier Characterization: They prove that a static ensemble lies exactly on the Pareto frontier of the trade-off between acceptance rate and distributional bias (Total Variation distance). This implies that a single static weight can theoretically achieve any point on this frontier, but it cannot move off the frontier to find a better balance for specific contexts.
• Dynamic Superiority: DIVERSED is shown to surpass the static Pareto frontier by adapting weights dynamically, achieving higher acceptance rates for the same quality, or better quality for the same acceptance rate.

Empirical Contributions

• State-of-the-Art Performance: DIVERSED consistently outperforms standard SD, lossy SD, SpecCascade, and static ensemble methods across multiple model pairs (Llama, Qwen, Gemma) and tasks (Math, Summarization, Coding).
• Cross-Task Generalization: The method can be trained on one dataset (e.g., GSM8K) and effectively applied to another (e.g., CNNDM), though performance is best when trained on the target task.
• Efficiency-Quality Trade-off: DIVERSED achieves a superior Pareto frontier, offering lower wall-clock latency without sacrificing generation quality (ROUGE-2, Accuracy, Pass@1).

4. Experimental Results

The paper evaluates DIVERSED on four benchmarks: GSM8K (Math), CNNDM & XSum (Summarization), and MBPP (Coding).

• Acceptance Rates: DIVERSED achieves significantly higher acceptance rates than all baselines.
  • Example (Llama 8B/1B on GSM8K): DIVERSED achieved 72.61% acceptance vs. 69.49% for Static Ensemble and 44.60% for standard SD.
  • Example (Qwen 8B/0.6B on GSM8K): DIVERSED achieved 76.48% vs. 67.58% for Static Ensemble.
• Generation Quality: Despite higher acceptance rates, DIVERSED maintains generation quality comparable to or better than the target model.
  • On GSM8K, accuracy remained stable (e.g., 67% for Llama pair) while speed increased.
  • On CNNDM, ROUGE-2 scores improved (e.g., 12.11 for DIVERSED vs. 11.46 for Static Ensemble).
• Speedup: The increase in acceptance rate directly correlates to reduced wall-clock time. DIVERSED consistently accepted more tokens per verification round (e.g., ~3.8 tokens vs. ~2.9 for baselines), leading to substantial latency reductions.
• Ablation: Fine-tuning the draft model alone did not reliably improve acceptance rates compared to the dynamic ensemble approach, suggesting that distributional alignment via the verifier is more critical than the draft model's standalone performance.

5. Significance

• Breaking the Rigid Constraint: DIVERSED challenges the dogma that speculative decoding must strictly match the target distribution to preserve quality. It demonstrates that "relaxed" verification, when done intelligently and dynamically, can be both faster and more accurate.
• Context-Aware Verification: By learning to weigh the draft and target models based on the specific context, DIVERSED solves the "one-size-fits-all" problem of static relaxation. It knows when to trust the draft model and when to defer to the target.
• Practical Deployment: The method requires no architectural changes to the target or draft models (unlike EAGLE or Medusa) and only adds a lightweight ensemble head, making it highly compatible with existing LLM infrastructure.
• Theoretical Rigor: The derivation of the expected acceptance length without i.i.d. assumptions provides a more accurate theoretical foundation for analyzing speculative decoding efficiency.

In summary, DIVERSED represents a significant step forward in LLM inference acceleration, offering a learnable, dynamic mechanism to maximize throughput while strictly controlling generation quality.