S2D: Sorted Speculative Decoding For More Efficient Deployment of Nested Large Language Models

TL;DR

This work tackles the high costs of deploying autoregressive LLMs by introducing Sorted Speculative Decoding (S2D), a framework that trains multiple draft sub-models inside a single draft via Sorted Fine-tuning (SoFT) and employs an adaptive draft selection strategy to serve multiple target models without modifying them. By leveraging a multi-target draft architecture and confidence-based draft generation, S2D achieves substantial speedups (averaging around 1.55x) across targets from 7B to 70B parameters and consistently outperforms vanilla speculative decoding baselines on Spec-Bench benchmarks. The approach reduces deployment complexity and expense while maintaining accuracy, and ablation studies show the impact of thresholds, attention trees, feature alignment, and pre-training on performance. Overall, S2D enables efficient, scalable deployment of diverse LLMs by reusing a single draft across multiple targets with tunable trade-offs between draft latency and accepted token length.

Abstract

Deployment of autoregressive large language models (LLMs) is costly, and as these models increase in size, the associated costs will become even more considerable. Consequently, different methods have been proposed to accelerate the token generation process and reduce costs. Speculative decoding (SD) is among the most promising approaches to speed up the LLM decoding process by verifying multiple tokens in parallel and using an auxiliary smaller draft model to generate the possible tokens. In SD, usually, one draft model is used to serve a specific target model; however, in practice, LLMs are diverse, and we might need to deal with many target models or more than one target model simultaneously. In this scenario, it is not clear which draft model should be used for which target model, and searching among different draft models or training customized draft models can further increase deployment costs. In this paper, we first introduce a novel multi-target scenario for the deployment of draft models for faster inference. Then, we present a novel, more efficient sorted speculative decoding mechanism that outperforms regular baselines in multi-target settings. We evaluated our method on Spec-Bench in different settings, including base models such as Vicuna 7B, 13B, and LLama Chat 70B. Our results suggest that our draft models perform better than baselines for multiple target models at the same time.

Figures (4)

• Figure 1: The figure on the left illustrates the draft model training process, comparing supervised fine-tuning (SFT) with Sorted fine-tuning (SoFT) using two sub-models with 6 and 9 layers. The figure on the right demonstrates the confidence-based drafting process, where the SoFT draft model is utilized to generate candidate tokens. The confidence thresholds for the two sub-models are set at 0.75 and 0.7, respectively.
• Figure 2: Comparison among Speedup ratios of speculative and S2D methods on different domains on multiple targets.
• Figure 3: Speedup ratios based on difference confidence thresholds in multiple targets on MT-Bench. In the thresholds axis labels, the first, second and third numbers represent the thresholds for the first draft sub-model (layer 6), the second draft sub-model (layer 9) and the last one (layer 12). The temperature was set to 0.0 in these experiments.
• Figure 4: Comparison among Speedup ratios of speculative and S2D methods on different domains on multiple targets.