Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Balancing Coverage and Draft Latency in Vocabulary Trimming for Faster Speculative Decoding

Ofir Ben Shoham

TL;DR

A utility function with a Tree-structured Parzen Estimator is optimized to efficiently explore the coverage-latency Pareto frontier under a minimum coverage constraint and casts draft vocabulary selection as a constrained optimization problem that balances token coverage and draft latency.

Abstract

Speculative decoding accelerates inference for Large Language Models by using a lightweight draft model to propose candidate tokens that are verified in parallel by a larger target model. Prior work shows that the draft model often dominates speculative decoding latency, since it generates tokens sequentially and incurs high cost from its language modeling head as vocabulary size grows. This exposes a fundamental trade-off in draft model design: larger vocabularies improve token coverage and agreement with the target model, but incur higher draft latency, while smaller vocabularies reduce latency at the risk of missing tokens required for accurate draft generation. We address this trade-off through vocabulary trimming for draft models, motivated by the observation that domain-specific workloads use only a small fraction of the full vocabulary. We cast draft vocabulary selection as a constrained optimization problem that balances token coverage and draft latency. Coverage is computed over assistant responses in the training data, while latency is estimated using architecture-aware FLOPs that capture the cost of the language modeling head as a function of vocabulary size. We optimize a utility function with a Tree-structured Parzen Estimator to efficiently explore the coverage-latency Pareto frontier under a minimum coverage constraint. Experiments show improved speculative decoding throughput while reducing draft vocabularies by up to 97% with high coverage. On domain-specific tasks, we achieve up to 16% latency reduction and 20% throughput improvement, and up to 6.7% throughput gains on diverse out-of-distribution tasks.

Balancing Coverage and Draft Latency in Vocabulary Trimming for Faster Speculative Decoding

TL;DR

A utility function with a Tree-structured Parzen Estimator is optimized to efficiently explore the coverage-latency Pareto frontier under a minimum coverage constraint and casts draft vocabulary selection as a constrained optimization problem that balances token coverage and draft latency.

Abstract

Speculative decoding accelerates inference for Large Language Models by using a lightweight draft model to propose candidate tokens that are verified in parallel by a larger target model. Prior work shows that the draft model often dominates speculative decoding latency, since it generates tokens sequentially and incurs high cost from its language modeling head as vocabulary size grows. This exposes a fundamental trade-off in draft model design: larger vocabularies improve token coverage and agreement with the target model, but incur higher draft latency, while smaller vocabularies reduce latency at the risk of missing tokens required for accurate draft generation. We address this trade-off through vocabulary trimming for draft models, motivated by the observation that domain-specific workloads use only a small fraction of the full vocabulary. We cast draft vocabulary selection as a constrained optimization problem that balances token coverage and draft latency. Coverage is computed over assistant responses in the training data, while latency is estimated using architecture-aware FLOPs that capture the cost of the language modeling head as a function of vocabulary size. We optimize a utility function with a Tree-structured Parzen Estimator to efficiently explore the coverage-latency Pareto frontier under a minimum coverage constraint. Experiments show improved speculative decoding throughput while reducing draft vocabularies by up to 97% with high coverage. On domain-specific tasks, we achieve up to 16% latency reduction and 20% throughput improvement, and up to 6.7% throughput gains on diverse out-of-distribution tasks.
Paper Structure (21 sections, 5 equations, 3 figures, 6 tables)

This paper contains 21 sections, 5 equations, 3 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Our Contributions:
  3. Method
  4. Problem Formulation
  5. Token Coverage Estimation
  6. Draft Model Latency Estimation
  7. Utility Function Definition
  8. TPE-Based Optimization
  9. Draft Vocabulary Size Optimization
  10. Experiments
  11. Results on out of distribution datasets
  12. In-domain Results
  13. Discussion
  14. Generalization to Out-of-Distribution Datasets
  15. Token Coverage Analysis.
  16. ...and 6 more sections

Figures (3)

  • Figure 1: Utility score vs. draft vocabulary size. The utility function peaks at 13,264 tokens, representing the optimal tradeoff between coverage and latency while training on the Open-PerfectBlend dataset xu2024perfect.
  • Figure 2: Pareto front showing the coverage-latency tradeoff of the draft model, while training on the Open-PerfectBlend dataset xu2024perfect. The optimal point (red star) balances high coverage with substantial latency reduction.
  • Figure 3: Optimal vocabulary size based on our approach vs. training data size. The optimal vocabulary converges to approximately 13K tokens after 10K samples, demonstrating stability across different random subsets of the training data.