BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

Technical Summary: BudgetDraft – Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

1. Problem Statement

The paper addresses a critical bottleneck in deploying Speculative Decoding (SD) for mid-to-long context inference (4K–16K tokens). While SD accelerates autoregressive (AR) generation by using a small "drafter" model to propose tokens for a larger "verifier" model to validate, standard deployment faces a sparse/full mismatch.

• The Constraint: To manage peak GPU memory (VRAM) and latency, the drafter is often forced to operate with a sparse KV cache under a fixed budget, while the verifier maintains a full KV cache to ensure output quality.
• The Failure Mode: As context length increases from 4K to 16K, the alignment between the sparse-drafter and full-verifier degrades rapidly. The paper identifies a "16K collapse," where acceptance rates drop to near-zero, rendering SD ineffective (speedup $\le$ 1$\times$).
• The Sensitivity Issue: Naive sparse drafting is highly sensitive to the specific KV budget. A drafter tuned for a single budget often fails when deployed under different memory constraints, and the relationship between budget size and acceptance rate is non-monotonic (smaller budgets sometimes yield higher acceptance than larger ones).

2. Methodology: BudgetDraft

The authors propose BudgetDraft, a training-time solution that creates a single, budget-robust drafter capable of maintaining stable acceptance rates across varying KV budgets without requiring extra inference-time components.

Core Architecture:
The method employs a Multi-View Sparse Training strategy where the drafter (student) is trained to align with a frozen verifier (teacher) under multiple sparse conditions simultaneously.

• Teacher Targets: The verifier generates greedy teacher targets ($x^*_t$) using its full KV cache.
• Dual-Branch Loss Function: The training objective combines two complementary losses:
  1. Acceptance-Aware Loss ($L_A$): A full-cache branch where the drafter predicts the teacher's greedy tokens using a full prefix cache. This ensures direct alignment with the verifier's acceptance criterion (top-1 match).
  2. Multi-View Sparse Loss ($L_C$): A sparse-cache branch where the drafter processes the same sequence but with a randomly sampled KV budget ($B$) from a set (e.g., $\{256, 512, 1024, 2048\}$). The sparse cache is constructed by chunking the prefix and retaining top chunks based on cumulative attention weights.
• Training Objective: The total loss is $L = L_A + \lambda L_C$. By exposing the drafter to diverse sparse views of the same prefix during training, it learns to generalize across sparsity levels rather than overfitting to a single budget.
• Inference: At deployment, the trained drafter uses a sparse KV cache with a chosen budget $B$ to draft tokens. The verifier validates them using its full cache. No structural changes or intermediate models are added to the inference pipeline.

3. Key Contributions

1. Characterization of Acceptance Collapse: The paper empirically establishes that naive sparse speculative decoding suffers a severe acceptance rate collapse as context length moves from 4K to 16K, identifying 16K as a practical failure boundary. It also reveals the non-monotonic relationship between KV budget size and acceptance rate.
2. BudgetDraft Framework: The proposal of a multi-view training method that combines acceptance-aware alignment with multi-budget sparse training. This produces a budget-invariant drafter that recovers acceptance rates across all sparsity levels without extra inference-time overhead.
3. Empirical Validation: Demonstration that BudgetDraft achieves significant end-to-end speedups over standard AR decoding while maintaining memory efficiency, specifically in the challenging 8K–16K regime.

4. Experimental Results

Experiments were conducted on three datasets (PG-19, LongBench, LWM) using a 68M parameter drafter and a 7B parameter verifier on a single NVIDIA A100 GPU.

• Speedup Performance:
  • 4K Context: Achieved up to 6.55$\times$ speedup over AR.
  • 8K Context: Achieved up to 4.46$\times$ speedup over AR.
  • 16K Context: Achieved up to 2.10$\times$ speedup over AR (where naive SD often fails completely).
• Robustness: Unlike baseline SD (sparse/full), which shows sharp performance drops or near-zero acceptance as context length increases or budgets change, BudgetDraft maintains stable acceptance rates across different KV budgets ($B \in \{256, 512, 1024, 2048\}$).
• Comparisons: BudgetDraft outperformed structural mitigation methods like TriForce and advanced drafting methods like EAGLE-3 in the 8K–16K regime. For instance, at 16K on LWM, BudgetDraft achieved ~1.94$\times$ speedup, compared to 1.19$\times$ for TriForce and 1.36$\times$ for EAGLE-3.
• Ablation: Removing the multi-view sparse loss ($L_C$) resulted in unstable performance under budget shifts, confirming the necessity of multi-budget sampling during training.

5. Significance and Claims

The paper claims that BudgetDraft solves the "sparse/full mismatch" problem that currently limits the practical deployment of speculative decoding in long-context scenarios. By shifting the complexity to the training phase, it enables a simple, memory-friendly inference pipeline that is robust to varying hardware constraints (VRAM availability) and workload concurrency.

The authors emphasize that their approach does not require complex structural changes (like retrieval caches) or multiple models at inference time. Instead, it provides a single, lightweight drafter that can adapt to different deployment budgets, making speculative decoding viable for mid-to-long context applications (4K–16K) where it was previously ineffective.

Limitations Acknowledged by the Paper:

• Positional Extrapolation: The 68M drafter has a native position embedding limit (2048), requiring extrapolation for 8K–16K contexts, which may still limit acceptance rates.
• Verifier Specificity: The drafter is trained against a specific verifier; retraining is required if the verifier model changes.
• Greedy Decoding: Results are based on greedy decoding; extending to sampling-based decoding (e.g., temperature sampling) is left for future work.