Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TokenTiming: A Dynamic Alignment Method for Universal Speculative Decoding Model Pairs

Sibo Xiao, Jinyuan Fu, Zhongle Xie, Lidan Shou

TL;DR

TokenTiming introduces a universal speculative decoding framework that removes the shared-vocabulary constraint between draft and target models by employing Dynamic Token Warping to construct a lossless mapping between heterogeneous vocabularies. The method converts draft tokens into proxy target tokens and uses a DTW path to transfer the draft distribution to the target vocabulary, enabling on-the-fly alignment without retraining. Empirical results across translation, summarization, code, and math demonstrate speedups up to $1.57\times$ and competitive performance close to homogeneous-vocabulary SD, with robust results across diverse draft models. By enabling plug-and-play draft models and model switches, TokenTiming broadens the practical deployment of speculative decoding despite some overhead and sensitivity to tokenization granularity. Overall, the approach broadens SD applicability and offers a practical route to faster LLM inference at real-world scales.

Abstract

Accelerating the inference of large language models (LLMs) has been a critical challenge in generative AI. Speculative decoding (SD) substantially improves LLM inference efficiency. However, its utility is limited by a fundamental constraint: the draft and target models must share the same vocabulary, thus limiting the herd of available draft models and often necessitating the training of a new model from scratch. Inspired by Dynamic Time Warping (DTW), a classic algorithm for aligning time series, we propose the algorithm TokenTiming for universal speculative decoding. It operates by re-encoding the draft token sequence to get a new target token sequence, and then uses DTW to build a mapping to transfer the probability distributions for speculative sampling. Benefiting from this, our method accommodates mismatched vocabularies and works with any off-the-shelf models without retraining and modification. We conduct comprehensive experiments on various tasks, demonstrating 1.57x speedup. This work enables a universal approach for draft model selection, making SD a more versatile and practical tool for LLM acceleration.

TokenTiming: A Dynamic Alignment Method for Universal Speculative Decoding Model Pairs

TL;DR

TokenTiming introduces a universal speculative decoding framework that removes the shared-vocabulary constraint between draft and target models by employing Dynamic Token Warping to construct a lossless mapping between heterogeneous vocabularies. The method converts draft tokens into proxy target tokens and uses a DTW path to transfer the draft distribution to the target vocabulary, enabling on-the-fly alignment without retraining. Empirical results across translation, summarization, code, and math demonstrate speedups up to and competitive performance close to homogeneous-vocabulary SD, with robust results across diverse draft models. By enabling plug-and-play draft models and model switches, TokenTiming broadens the practical deployment of speculative decoding despite some overhead and sensitivity to tokenization granularity. Overall, the approach broadens SD applicability and offers a practical route to faster LLM inference at real-world scales.

Abstract

Accelerating the inference of large language models (LLMs) has been a critical challenge in generative AI. Speculative decoding (SD) substantially improves LLM inference efficiency. However, its utility is limited by a fundamental constraint: the draft and target models must share the same vocabulary, thus limiting the herd of available draft models and often necessitating the training of a new model from scratch. Inspired by Dynamic Time Warping (DTW), a classic algorithm for aligning time series, we propose the algorithm TokenTiming for universal speculative decoding. It operates by re-encoding the draft token sequence to get a new target token sequence, and then uses DTW to build a mapping to transfer the probability distributions for speculative sampling. Benefiting from this, our method accommodates mismatched vocabularies and works with any off-the-shelf models without retraining and modification. We conduct comprehensive experiments on various tasks, demonstrating 1.57x speedup. This work enables a universal approach for draft model selection, making SD a more versatile and practical tool for LLM acceleration.
Paper Structure (67 sections, 1 theorem, 13 equations, 14 figures, 11 tables, 5 algorithms)

This paper contains 67 sections, 1 theorem, 13 equations, 14 figures, 11 tables, 5 algorithms.

Table of Contents

  1. Introduction
  2. Limited Selection of Draft Models.
  3. Costly and Inflexible Training.
  4. Related Works
  5. SD with Homogeneous Vocabularies
  6. SD with Pruned Vocabularies
  7. SD with Heterogeneous Vocabularies
  8. Preliminaries
  9. Speculative Decoding (SD)
  10. Standard Speculative Sampling
  11. Vocabulary Mismatch
  12. Method
  13. Methodological Core
  14. Methodological Pipeline
  15. Draft Token Conversion
  16. ...and 52 more sections

Key Result

Lemma 1

The proxy sequence $T = (t_1, \dots, t_m) = f_{\text{encode}}^t(f_{\text{decode}}^d(D))$ is identical to the draft sequence $D$ in the string space.

Figures (14)

  • Figure 1: Comparison of core points of TokenTiming with previous work. Previous methods (TLI) carries the transfer of probability distribution in the vocabulary space of $\mathcal{V}_d \cap \mathcal{V}_t$. If the draft tokens fall outside the intersection, they will be accepted unconditionally, which violates the integrity of losslessness. TokenTiming constructs intact token mapping via DTW, ensuring lossless speculative sampling.
  • Figure 2: Phase illustrations of TokenTiming. (a) illustrates the re-tokenization of Draft Tokens into Proxy Target Tokens, which are used to construct the mapping in the DTW. (b) DTW calculation process (token distance matrix) and aligned token mapping $\pi ^* =[(S,Scale),(cal,Scale),(ing,ing),(Law,L),(Law,aw)]$. The calculation rules for this mapping are presented in Alg. \ref{['alg:dtw_core']}. (c) Probability distribution of draft tokens is transferred based on the mapping $\pi^*$. (d) Dynamic alignment in every Draft-Verify iteration of Speculative Decoding.
  • Figure 3: Total Time Cost (with TokenTiming overhead)
  • Figure 4: Speed-up vs. homogeneous-vocabulary SOTA on various target model scales and draft models. TokenTiming bridges the performance gap while keeping the universal draft-model advantage.
  • Figure 5: Performance metrics under various settings of $w$: $w=4$, $w=8$, $w=16$, and $w=\infty$.
  • ...and 9 more figures

Theorems & Definitions (4)

  • Lemma 1: String-level Consistency
  • proof
  • proof
  • proof