PaperlayerView exact revision on arXiv

Source-linked AI summary

T2D: Target-to-Draft Model Pruning & Adaptation for Efficient Speculative Decoding

Shahriar Kabir Nahin, Xiao Ye, Muhao Chen, Ben Zhou, Anshuman Chhabra

arXiv:2610.34103v1cs.CLcs.AIcs.LG

TL;DR

Existing draft-model construction does not jointly account for latency and alignment with the target, although both shape speculative-decoding efficiency. T2D uses latency-aware structured pruning and distillation to learn a target-derived draft; experiments report better speedups than baselines across domains and model families.

  • Problem

    Draft construction often uses independently trained models or manually chosen pruned architectures, without explicitly balancing target alignment and decoding efficiency.

    Source: §1 · p. 2

  • Method

    T2D constructs a draft from the target through latency-aware structured pruning and distillation, jointly optimizing architecture and target alignment.

    Source: §4 · p. 6, §6 · p. 10

  • Results

    T2D outperforms baselines across all domains and achieves 1.89× overall speedup for Vicuna-13B-v1.5.

    Source: §5.1 · p. 8

  • Takeaways & Limitations

    T2D-derived drafts can be further adapted with AdaSpec, where they achieve up to 20% higher speedup than other baselines on GSM8K.

    Source: §5.4 · p. 10

Abstract

from arXiv · show

Speculative decoding significantly accelerates autoregressive text generation in Large Language Models (LLMs), but its efficacy depends largely on two key factors: the distributional alignment between the draft and the target model, and the draft model's latency during inference. However, existing approaches typically utilize independently pretrained small draft models or seek to prune the target model to a manually chosen architecture, thereby failing to explicitly balance the trade-off between draft model alignment and token generation efficiency. To address this gap, we formulate draft model design as an optimization problem and propose T2D, a framework that constructs draft models directly from the target model using latency-aware structured pruning and knowledge distillation. In this manner, T2D obtains draft models that are jointly optimized for both measured inference latency and cross-entropy loss on the target model's responses, resulting in superlative speculative decoding performance. Through extensive experiments on multi-domain speculative decoding benchmarks and across several target model families of differing scale, we demonstrate how T2D prunes up to 92\% of the target model to obtain highly performant draft models compared to competitive baselines. Moreover, draft models obtained via T2D can also be utilized by other speculative decoding token-alignment and knowledge distillation methods in a plug-and-play fashion, leading to further gains. Our findings underscore the importance of jointly navigating the draft model architecture-alignment frontier to unlock the full potential of speculative decoding for LLMs.

1 INTRODUCTION

Speculative decoding trades draft-model latency against target-draft agreement, and existing construction methods do not jointly optimize those factors. T2D learns a pruned draft from the target while distilling its knowledge, with reported gains over pruning-based baselines.

1 INTRODUCTION

  • A low-latency draft with frequent rejections or an accurate but computationally expensive draft can both undermine speculative-decoding efficiency.Figure 1 illustrates the trade-off using traces for the same output sequence with different draft latencies and acceptance ratios.Source: §1 · p. 2, Figure 1 · p. 2
  • Prior approaches either adapt independently trained small models or prune the target with unstructured sparsity, which is difficult to exploit with conventional dense GPU kernels.Structured pruning can remove architectural components, but prior methods manually choose the draft architecture before considering efficiency and alignment together.Source: §1 · p. 2
  • T2D jointly optimizes draft architecture and target-model alignment, allowing the architecture to emerge during latency-aware pruning and distillation.The introduction reports up to 1.89× overall speculative-decoding speedup across 11 domains and up to 2.46× in multilingual understanding.Source: §1 · pp. 2–3
  • T2D constructs drafts directly from the target and reports outperforming ShortGPT-style pruning and fixed-architecture joint pruning-and-distillation baselines.The reported evaluations cover Llama-2, Vicuna-v1.5, Qwen2.5, Qwen3, and Nemotron model families.Source: §1 · p. 3

2 RELATED WORKS

Prior work changes speculative drafting or verification algorithms and constructs draft models through separate training, target-derived pruning, or structural pruning. T2D is motivated by searching structural choices for speculative-decoding performance rather than fixing them in advance.

Speculative Decoding

Draft Model Construction

  • Draft construction ranges from separately pretrained models to target-derived unstructured pruning and structural pruning that produces compact dense models.SD2 (Lasby et al., 2025) combines target-model unstructured pruning with self-distillation; Wei et al. (2026) remove Transformer layers, while Yan et al. (2025) study depth-width configurations.Source: §2 · p. 3

Structured Pruning

  • Structured pruning removes layers, attention heads, or feed-forward channels, and Sheared-LLaMA uses learnable gates under predefined architectural constraints.The section notes that many such methods evaluate standard response-quality benchmarks rather than speculative-decoding performance.Source: §2 · pp. 3–4
  • The authors report that prior work has not explored speculative-decoding-aware pruning as a learnable objective for improving speculative-decoding performance.This motivates searching structural choices rather than fixing the draft architecture beforehand.Source: §2 · pp. 3–4

3 PRELIMINARIES AND BACKGROUND

The background explains how model architecture affects decoding cost and how speculative decoding uses a draft model verified by a target model. It frames draft construction as balancing latency reduction with target-draft alignment through pruning and distillation.

LLMs and Autoregressive Decoding

  • An autoregressive LLM predicts each next token from the prompt and prior tokens, while depth, hidden dimension, attention heads, and feed-forward width shape capacity and decoding cost.A larger layer count is associated with greater estimated single-step latency.Source: §3 · p. 4

Speculative Decoding

  • Speculative decoding has a draft propose a sequence that the target verifies in one pass, retaining only the initial consecutive matching tokens and using the target prediction at the first mismatch.Each verification round extends the sequence by the accepted draft tokens plus one correction token.Source: §3 · p. 4
  • For a fixed draft length, speedup benefits from lower draft-to-target latency ratio and more accepted draft tokens.The background identifies draft latency, target verification latency, average accepted tokens, and draft length as primary speedup factors.Source: §3 · p. 4

T2D construction

  • T2D’s construction algorithm takes a target model, its latency, alignment data, a latency predictor, a reference cross-entropy value, and an epoch count as inputs.The listed procedure optimizes gates, selects retained components, instantiates a dense draft, and fine-tunes it on the alignment data.Source: §3 · p. 5

Knowledge Distillation for Draft-Model Alignment

  • Knowledge distillation aligns the draft with the target, including by training on target-generated sequences with causal cross-entropy.The excerpt also identifies direct KL-divergence minimization between next-token distributions as a distillation approach.Source: §3 · p. 5

Structured Pruning

  • Structured pruning removes architectural components to produce a smaller dense model, but conventional pipelines generally specify its architecture beforehand.T2D instead aims to prune while maintaining target-model alignment, learning the efficiency-alignment trade-off for draft construction.Source: §3 · p. 5

4 PROPOSED METHOD: T2D

T2D formulates draft construction as jointly optimizing architecture and target alignment to balance decoding latency with alignment. It uses differentiable structured pruning and a joint loss to derive and fine-tune a compact draft model.

4 PROPOSED METHOD: T2D

  • T2D frames draft construction as selecting a structured-pruned architecture that balances decoding latency and alignment with the target model.The search space comprises architectures reachable through structured pruning, evaluated on query–response pairs generated by the target model.Source: §4 · pp. 5–6
  • Learnable gates with Hard-Concrete relaxation represent removable Transformer structures as differentiable keep-or-remove decisions.The gates cover layers, attention structures, and feed-forward channels; expected-open probabilities estimate the resulting architecture dimensions.Source: §4 · p. 6
  • A frozen latency predictor estimates latency from architecture dimensions, allowing latency gradients to reach the pruning gates.The predictor is trained offline using measured single-token decoding latency for candidate dense architectures.Source: §4 · p. 6
  • The joint objective combines normalized predicted latency, target-response cross-entropy, and a reference-guided penalty for excessive CE degradation.The penalty activates when smoothed CE exceeds its reference, while the shared gate parameters optimize latency and alignment together.Source: Figure 2 · p. 6, §4 · p. 7
  • After joint training, T2D converts gate scores into discrete structure choices, copies selected target parameters, and fine-tunes the dense draft for one epoch.The fine-tuning is intended to let weights readjust after discretization.Source: §4 · p. 7

5 EXPERIMENTS AND RESULTS

Experiments compare T2D with pruning-based draft-model baselines across target models, domains, and alignment settings. T2D achieves the strongest reported speedups while balancing latency and alignment, and its drafts can also be adapted to AdaSpec.

5 EXPERIMENTS AND RESULTS

  • Experiments cover five target models, 250,000 training and 5,000 validation pairs, and SPEED Bench’s 798 queries across 11 domains, measuring speedup against autoregressive decoding.The targets span Llama-2, Vicuna, Qwen2.5, Qwen3, and Nemotron; Qwen3 and Nemotron are reasoning models.Source: §5 · pp. 7–8

5.1 MAIN RESULTS

  • Across domains, T2D leads all baselines; Llama-2-7B-Chat reaches 1.74× overall speedup, 20% above ShortGPT and 9.4% above Wide.Vicuna-13B-v1.5 reaches 1.89× overall speedup, more than 6.2% above heuristic baselines, and 2.46× in Multilingual. T2D also produces coherent Llama-2 responses after 88% pruning, while heuristic drafts sometimes fall below 1.00×.Source: Table 1 · p. 8, §5.1 · pp. 8–9

5.2 WHY IS T2D BETTER THAN RANDOM MANUAL ARCHITECTURE SELECTION?

  • In ShortGPT’s Llama-2-7B-Chat sweep, speedup peaks at 89% pruning, while T2D reaches 88% and automates the latency–acceptance trade-off.Beyond the peak, speedup declines; other pruning baselines require a fixed prior pruning budget.Source: Figure 3 · p. 9, §5.2 · p. 9

5.3 HOW DOES T2D PRUNE THE MODEL WHILE BALANCING CE LOSS AND LATENCY?

  • On Llama-2-7B-Chat, pruning starts near 10% of training, then CE loss is rebalanced against latency until pruning plateaus at 88%.Layer pruning supplies most latency benefit; larger models also prune heads and FFN as those components grow in dimension and latency contribution.Source: §5.3 · p. 9, Figure 4 · p. 9

5.4 CAN T2D MODELS ADAPT TO OTHER ALIGNMENT METHODS?

  • Under GSM8K AdaSpec distillation, T2D outperforms the other baselines by up to 20% in speedup.The result supports using T2D-derived drafts in a plug-and-play fashion with another alignment method.Source: Table 2 · p. 10, §5.4 · p. 10

5.5 OUR T2D DRAFT MODELS BETTER THAN AN INDEPENDENTLY TRAINED DRAFT MODEL?

  • T2D’s draft reaches 1.74× overall speedup versus 1.18× for Sheared-Llama-1.3B-ft, despite the latter’s 340× more pretraining data.Sheared-Llama-1.3B-ft was fine-tuned for one epoch on Llama-2-7B-Chat-generated responses.Source: §5.5 · p. 10

6 CONCLUSION

T2D jointly optimizes draft architecture and target-model alignment through latency-aware structured pruning and distillation. Across five model families, it outperforms heuristic architecture choices and pruning-based baselines.

6 CONCLUSION

  • T2D consistently outperforms heuristic architecture choices and pruning-based baselines across the Llama-2, Vicuna-v1.5, Qwen2.5, Qwen3, and Nemotron families.Table 3 compares speculative-decoding speedup for T2D and finetuned Sheared-Llama using Llama-2-7B-Chat as the target.Source: §6 · p. 10, Table 3 · p. 10
  • T2D jointly optimizes draft-model architecture and target-model alignment, allowing the draft architecture to emerge from the optimization process.Source: §6 · p. 10
  • Its pruning objective combines a differentiable inference-latency estimate with target-distilled cross-entropy and a reference-based penalty that prevents pruning from outpacing alignment.Source: §6 · p. 10

A BENCHMARK DATASET DETAILS

The evaluation uses 798 SPEED Bench samples spanning 11 domains, with a standard speculative-decoding draft length of four tokens.

A BENCHMARK DATASET DETAILS

B TRAINING DATASET

Training data comprises 250K instruction-tuning queries used to elicit target-model responses; the queries cover about 95% of the Llama tokenizer vocabulary.

B TRAINING DATASET

  • 250K randomly sampled instruction-tuning queries cover approximately 95% of the Llama tokenizer vocabulary and correspond to approximately 0.022B Llama tokens.Source: B TRAINING DATASET · p. 15, Table 5 · p. 15
  • The same query set is used across target models, but generated response counts and average response lengths vary by model.Table 5 reports total generated response tokens and average response length for each target model.Source: B TRAINING DATASET · p. 15, Table 5 · p. 15
  • The queries are sampled from open-source instruction-tuning datasets, and responses are generated using the target models.Source: B TRAINING DATASET · p. 15

C IMPLEMENTATION DETAILS

The implementation compares T2D with Wide and ShortGPT-based draft construction, using shared architecture and training settings across target models. Hyperparameters are summarized for Llama-2-7B-Chat, and experiments run on a 14-GPU DGX B200 server.

C IMPLEMENTATION DETAILS

  • Wide uses Yan et al.'s best-performing architecture, prunes each target model to that architecture, and receives supervised fine-tuning on the same dataset as T2D.Source: C IMPLEMENTATION DETAILS · p. 15
  • The experiments also compare a ShortGPT-based draft method with the same architecture as Wide.Both approaches use the same 5-layer Wide architecture across target models; Wide models are optimized for two epochs and fine-tuned for one further epoch after pruning.Source: C IMPLEMENTATION DETAILS · p. 16
  • Llama-2-7B-Chat hyperparameters are listed in Table 6; parameters for other models are similar and recommended for tuning to their architectures.Source: C IMPLEMENTATION DETAILS · p. 16, Table 6 · p. 16
  • Experiments ran on a Linux server with 14 NVIDIA DGX B200 GPUs, each with 192 GB VRAM.Source: C IMPLEMENTATION DETAILS · p. 16

D T2D-DERIVED DRAFT MODEL ARCHITECTURES

Table 7 catalogs the final draft-model architectures derived by T2D for all T2D models.

D T2D-DERIVED DRAFT MODEL ARCHITECTURES

E RATIO OF WEIGHTS, CE LOSS, AND LATENCY OVER THE TRAINING PERIOD FOR OTHER

T2D adjusts component pruning during optimization, achieves higher speedup than Sheared-Llama-1.3B in every domain, and uses a learned latency predictor to guide optimization.

E RATIO OF WEIGHTS, CE LOSS, AND LATENCY OVER THE TRAINING PERIOD FOR OTHER

G DEVELOPMENT OF THE MLP LATENCY PREDICTOR

  • Candidate architectures vary layer, head, and FFN dimensions; median repeated single-token latency measurements on target hardware provide regression targets.The predictor uses estimated retained layer, head, and FFN dimensions to predict cached decoding-step log latency.Source: G DEVELOPMENT OF THE MLP LATENCY… · p. 18
  • A three-layer MLP predicts log latency using Huber loss and input perturbations, then remains frozen to provide differentiable estimates during joint optimization.The perturbations encourage smooth interpolation between measured architectures.Source: G DEVELOPMENT OF THE MLP LATENCY… · p. 18

F PER DOMAIN SPEEDUP COMPARISON WITH SHEARED LLAMA

H QUALITATIVE ANALYSIS OF SPECULATIVE DECODING PERFORMANCE

T2D draft models generally produce relevant, coherent responses despite pruning of up to 92% of the target model. Across decoding traces, they share some early errors with baselines but also capture finer token details.

H QUALITATIVE ANALYSIS OF SPECULATIVE DECODING PERFORMANCE

Related papers

Loading 2610.34103v1…