Automatic Task Detection and Heterogeneous LLM Speculative Decoding

TL;DR

TaskSpec addresses the challenge of optimizing speculative decoding for downstream tasks by automating task partitioning, constructing task-specific heterogeneous draft models via LoRA, and routing prompts online through a lightweight Mamba-based classifier. By collecting vanilla speculative decoding outputs, clustering them into task groups, and training task-aligned draft models, TaskSpec significantly improves draft acceptance rates and reduces end-to-end latency across multiple tasks (text generation, mathematics, translation, chemistry QA). Experimental results on LLaMA-2-13B as the target and LLaMA-68M as the draft show speedups up to 2.64× and substantial gains in accepted tokens, with a compact, efficient prompt classifier achieving high accuracy and low inference cost. The approach automates cross-domain transfer, enabling flexible, efficient domain-specific speculative decoding suitable for real-world LLM inference workloads.

Abstract

Speculative decoding, which combines a draft model with a target model, has emerged as an effective approach to accelerate large language model (LLM) inference. However, existing methods often face a trade-off between the acceptance rate and decoding speed in downstream tasks due to the limited capacity of the draft model, making it difficult to ensure efficiency across diverse tasks. To address this problem, we propose a speculative decoding algorithm tailored for downstream task optimization. It includes an automatic task partitioning and assigning method, which automatically categorizes downstream tasks into different sub-tasks and assigns them to a set of heterogeneous draft models. Each draft model is aligned with the target model using task-specific data, thereby enhancing the consistency of inference results. In addition, our proposed method incorporates an online lightweight prompt classifier to dynamically route prompts to the appropriate draft model. Experimental results demonstrate that the proposed method improves draft accuracy by 6% to 50% over vanilla speculative decoding, while achieving a speedup of 1.10x to 2.64x in LLM inference.

Figures (10)

• Figure 1: Workflow of speculative decoding. (a) Autoregressive decoding of the draft model. (b) Parallel verification of the target model.
• Figure 2: The word "reaction" tends to co-occur with different tokens in the domains of chemistry (a) and mathematics (b).
• Figure 3: The overview of TaskSpec.
• Figure 4: Data-driven automated partitioning of downstream tasks.
• Figure 5: Construction of heterogeneous draft models for downstream tasks.