Source-linked AI summary
SlimSpec: Low-Rank Draft LM-Head for Accelerated Speculative Decoding
Anton Plaksin, Sergei Krutikov, Sergei Skvortsov, Alexander Samarin
TL;DR
Speculative decoding is limited by the drafter’s expensive full-vocabulary LM-head, despite the drafter’s lightweight backbone. SlimSpec compresses the drafter’s hidden representation with a low-rank factorization while preserving full vocabulary support, achieving substantial LM-head acceleration with competitive acceptance and end-to-end gains.
Problem
The drafter’s full-vocabulary projection is a major speculative-decoding bottleneck, while vocabulary-reduction approaches add curation, routing, inference, or training complexity.
Method
SlimSpec is a drop-in low-rank LM-head architecture that factorizes the projection and compresses the drafter’s hidden representation instead of reducing the output vocabulary.
Results
4-5× LM-head acceleration with acceptance length close to the full-vocabulary baseline enabled SlimSpec to surpass other evaluated methods by 8-9% in end-to-end speedup.
Takeaways & Limitations
SlimSpec offers a full-vocabulary alternative with minimal training and inference-pipeline changes, combining lower LM-head cost with acceptance close to the baseline.
Takeaways & Limitations
Experiments use EAGLE-3 heads, vLLM 0.17.1, and NVIDIA H200 hardware, so transfer to other drafter families and stacks is only by analogy or may differ.
Abstract
from arXiv · showhide
Speculative decoding speeds up autoregressive generation in Large Language Models (LLMs) through a two-step procedure, where a lightweight draft model proposes tokens which the target model then verifies in a single forward pass. Although the drafter network is small in modern architectures, its LM-head still performs projection to a large vocabulary, becoming one of the major computational bottlenecks. In prior work this issue has been predominantly addressed via static or dynamic vocabulary truncation. Yet mitigating the bottleneck, these methods bring in extra complexity, such as special vocabulary curation, sophisticated inference-time logic or modifications of the training setup. In this paper, we propose SlimSpec, a low-rank parameterization of the drafter's LM-head that compresses the inner representation rather than the output, preserving full vocabulary support. We evaluate our method with EAGLE-3 drafter across three target models and diverse benchmarks in both latency- and throughput-bound inference regimes. SlimSpec achieves $4\text{-}5\times$ acceleration over the standard LM-head architecture while maintaining a competitive acceptance length, surpassing existing methods by up to $8\text{-}9\%$ of the end-to-end speedup. Our method requires minimal adjustments of training and inference pipelines. Combined with the aforementioned speedup improvements, it makes SlimSpec a strong alternative across wide variety of draft LM-head architectures.
1 Introduction
Speculative decoding reduces autoregressive inference cost, but the drafter’s full-vocabulary LM-head remains a major bottleneck. SlimSpec addresses this bottleneck with low-rank hidden-representation compression, preserving full vocabulary support while retaining acceptance quality and improving end-to-end speedup.
- Speculative decoding uses a lightweight drafter to propose multiple tokens that the target model verifies in parallel, preserving the target output distribution.
- Draft token-logit computation is a major bottleneck because the drafter must project every drafted position over a vocabulary often exceeding 10^5 tokens.
- Existing static and dynamic vocabulary-reduction methods lower projection cost but add vocabulary curation, bookkeeping, routing, top-k selection, or training complexity.
- SlimSpec replaces the standard LM-head with a low-rank factorized projection that compresses hidden representations rather than reducing candidate tokens, preserving the full vocabulary.
- The acceptance-cost framework relates LM-head acceleration to acceptance preservation, supporting comparisons among draft LM-head designs.
- 4-5× LM-head latency reduction accompanies acceptance length close to the full-vocabulary baseline, while end-to-end speedup surpasses other evaluated methods by 8-9%.The evaluation uses EAGLE-3 across three target models, diverse benchmarks, decoding temperatures, and serving regimes.
2 Related Work
Prior LM-head acceleration methods reduce computation mainly through static or dynamic vocabulary selection, while low-rank factorization has generally been developed outside speculative decoding. These approaches trade computational savings against vocabulary coverage, implementation complexity, and applicability to the drafter-target setting.
- Static truncation methods reduce drafter LM-head cost by selecting smaller vocabularies using corpus or target-generation frequency statistics.BCL additionally optimizes vocabulary size for token coverage and latency and trains the drafter with the selected vocabulary.
- Table 1 compares vocabulary pruning, top-k overhead, hyperparameters, and asymptotic LM-head FLOPs across acceleration methods.
- Static vocabulary truncation assigns zero probability to excluded tokens, which typically reduces acceptance quality.Probability redistribution can recover acceptance for pruned vocabularies but does not itself make LM-head computation cheaper.
- Dynamic selection methods route each context to a small active vocabulary subset, improving the speedup–acceptance trade-off through context adaptation.
- Dynamic selection adds top-k-style ranking, sorting, irregular indexing, and weight gathering, which can become GPU bottlenecks relative to dense matrix multiplication.
- Low-rank LM-head factorization in standard language modeling is not directly applicable to the speculative-decoding drafter-target setup because it is designed for single models.
3 Performance Model for Draft LM-Head Acceleration
The performance model separates drafting cost into backbone and LM-head components, then relates LM-head acceleration and acceptance preservation to end-to-end speculative-decoding speedup. The resulting trade-off depends on how much the LM-head contributes within the deployment configuration.
- Throughput structure: Drafting time decomposes into backbone computation and LM-head projection: Tdraft = Tbackbone + Thead.This decomposition isolates LM-head cost from the remaining draft-side computation.
- Throughput structure: 45%-60% of draft cost comes from the LM-head, making its vocabulary projection a major computational bottleneck.The standard projection has complexity O(V d), with V vocabulary size and d drafter hidden dimension.
- Acceptance–cost trade-off: End-to-end gains require LM-head savings to compensate for acceptance changes and the non-head costs of verification, overhead, and the backbone.The model defines Tnon-head as Toverhead + Tverify + Tbackbone and evaluates methods through acceptance preservation and LM-head acceleration.
- Acceptance–cost trade-off: The minimum acceptance preservation needed for speedup rises when the LM-head contributes little and falls when it dominates the pipeline.As κ approaches 0, any acceptance loss is fatal; greater LM-head dominance tolerates more acceptance loss.
- Deployment dependence: The trade-off parameter κ depends on the full deployment configuration, including target-model size, batch size, decoding temperature, and sampling overhead.These factors shift the relative weights of head and non-head costs across serving regimes.
4 SlimSpec
SlimSpec replaces the drafter’s full-vocabulary LM-head with a low-rank representation that preserves output support while reducing computation. Its design targets the acceptance and training-consistency limitations of vocabulary truncation without requiring substantial pipeline changes.
- 4.1 LM-head architecture: SlimSpec replaces the standard full-vocabulary projection with a low-rank factorization that compresses the hidden representation.The rank r is the chosen architectural hyperparameter and is smaller than the drafter hidden dimension d.
- 4.1 LM-head architecture: SlimSpec reduces LM-head cost from O(V d) to O(rd + V r) while preserving the full target vocabulary.Because V is much larger than d, the FLOP reduction is approximately linear in r.
- 4.1 LM-head architecture: Unlike vocabulary truncation, SlimSpec keeps all V token logits available by generating them through a thinner hidden representation.The rank controls both compressed-state width and head computation, with useful ranks including d/4, d/8, and d/16.
- 4.2 Advantages over vocabulary truncation: Truncation with KL training creates a train-test distribution mismatch, while SlimSpec preserves the full output support used by the target distribution.The mismatch arises because training uses a truncated, renormalized target but inference verifies against the full target distribution.
- 4.3 Simplicity: SlimSpec uses dense matrix multiplications and requires only small LM-head changes without altering the drafter backbone or training pipeline.It avoids token-statistics preprocessing, token-index mappings, routing, and top-k-selection logic.
- 4.2 Advantages over vocabulary truncation: Vocabulary truncation can impose an acceptance upper bound because tokens outside Vtr receive zero draft probability, whereas SlimSpec is not subject to this bound.Under greedy decoding, acceptance collapses to zero when the target argmax lies outside the truncated vocabulary.
5 Experimental Settings
The experiments compare SlimSpec with full-vocabulary, static-truncation, and dynamic-truncation LM-head baselines under matched drafter training and production-like serving conditions. Evaluation spans three target models, three task domains, decoding temperatures, and latency- and throughput-oriented batch sizes.
- Baselines: The study compares SlimSpec with Full Vocab, VocabTrim, FR-Spec, BCL, VocabTrim-T, and SpecVocab baselines.The baselines cover full-vocabulary projection, static vocabulary truncation, and dynamic vocabulary truncation.
- Target models and training: Experiments use Llama-3.1-8B-Instruct, GPT-OSS-20B, and Qwen3-30B-A3B-Instruct-2507 target models.Training data comprises 660K Infinity-Instruct-0625 prompts with responses generated by the corresponding target model.
- Drafter training: All drafters use the EAGLE-3 training pipeline with n = 6 speculative tokens, shared weights across positions, and standard KL-divergence loss.For each target model, the draft backbone is fixed so methods differ only in LM-head design.
- Evaluation: Evaluation covers MT-Bench, HumanEval, and GSM8K under greedy and stochastic decoding with batch sizes 1 and 64.These settings represent instruction following, code generation, mathematical reasoning, latency-critical serving, and high-throughput serving.
- Metrics: Throughput in tokens per second is the primary metric, while average acceptance length τ measures drafter quality independently of raw throughput.Each speedup and τ result averages five identical runs with different random seeds.
6 Evaluation Results
SlimSpec occupies a favorable acceptance-cost region, reducing LM-head cost while preserving acceptance close to the full-vocabulary baseline. It outperforms the evaluated static and dynamic vocabulary-truncation methods in the reported comparisons.
- Baseline comparisons: Static vocabulary truncation exhibits a cost-acceptance frontier: smaller vocabularies reduce ν but proportionally degrade ρτ.VocabTrim surpasses FR-Spec, while training-aware VocabTrim-T performs similarly to post-training VocabTrim at evaluated vocabulary sizes.
- Baseline comparisons: SlimSpec achieves the largest LM-head acceleration while keeping ρτ close to 1 in the plotted Llama-3.1-8B, temperature-0, batch-size-1 setting.The figure shows end-to-end speedup decomposition in the (ν, ρτ) plane, with shaded regions indicating no improvement over the full-vocabulary baseline.
- Acceptance-cost trade-off: Dynamic SpecVocab preserves full-vocabulary acceptance (ρτ ≈1) while reducing LM-head latency approximately 60% at ranks r =d/8 and r =d/16.
- Acceptance-cost trade-off: SlimSpec with rank r=d/8 achieves approximately 5× lower LM-head cost while maintaining acceptance quality ρτ =0.99.The paper adopts r=d/8 as the default configuration for the remaining evaluation.
- End-to-end speedup: SlimSpec improves averaged speedup by more than 8.5% over static truncation for Llama-3.1-8B and by 8.9% over SpecVocab on GPT-OSS-20B at batch size 64.For Qwen3-30B-A3B, the observed advantage is at most 1-2% because the target head accounts for a smaller fraction of total latency.
7 Conclusion
The paper evaluates draft LM-heads through the acceptance-cost trade-off, arguing that lower head cost improves end-to-end performance only when acceptance is preserved. SlimSpec compresses the drafter’s hidden representation while retaining full-vocabulary support and achieves the strongest reported overall trade-off.
- Conclusion: A cheaper draft LM-head does not necessarily improve end-to-end throughput when its cost reduction is offset by fewer accepted draft tokens.
- Conclusion: SlimSpec compresses the drafter’s hidden representation, preserves full-vocabulary support, and moves primarily along the LM-head cost axis.
- Conclusion: The rank-d/8 configuration achieves approximately 4-5× LM-head acceleration and the strongest throughput among the compared methods in the test settings.
- Future work: Future work could improve acceptance at fixed cost through acceptance-oriented training objectives or position-wise adaptivity.
8 Limitations
The evaluation has limitations in hyperparameter selection, drafter-family coverage, hardware and framework dependence, and baseline coverage.
- Limitations: The rank r is manually chosen, and the paper provides no automatic procedure for selecting it.
- Limitations: All experiments use EAGLE-3 draft heads, so conclusions for MEDUSA and Hydra transfer only by analogy.
- Limitations: Measured speedups depend on vLLM 0.17.1 and NVIDIA H200 hardware and may differ on other inference stacks.
- Limitations: CORAL and DynaSpec are omitted because reference implementations were unavailable, narrowing the study’s baseline coverage.
B Limitations discussion
The evaluation establishes scope boundaries for SlimSpec: its gains depend on deployment conditions, tested drafter configurations, model scales, workloads, and implementation choices.
- Deployment scope: SlimSpec’s gains depend on the deployment setup because LM-head acceleration helps only when saved head latency is substantial relative to the remaining pipeline.Serving stacks, GPU kernels, batch size, verification costs, sampling overhead, and scheduler behavior can change realized throughput improvement.
- Drafter scope: The empirical study evaluates SlimSpec only with EAGLE-3 auxiliary drafters and does not characterize alternative drafter backbones, dynamic draft trees, position-specific heads, or standalone draft models.Main experiments also use a fixed draft-backbone architecture with shared weights across draft positions.
- Evaluation scope: The evaluation covers three target models and three benchmarks but excludes larger models, long contexts, multilingual and domain-specific workloads, tool use, high temperatures, alternative sampling policies, and other serving frameworks.Reported speedups are measurements for the evaluated production-like setup rather than hardware-independent constants.
- Configuration scope: The best tested rank is r = d/8, but the paper provides no automatic rank-selection procedure because the optimal rank depends on model, vocabulary, serving, and benchmark factors.Determining the best rank therefore requires hyperparameter tuning in the tested settings.
- Comparison scope: The comparison omits CORAL and DynaSpec because faithful reproduction required additional machinery or unavailable reference implementations, and optimized dynamic-vocabulary kernels could change relative overheads.The evaluated baselines were selected for faithful implementation under a shared training and inference pipeline.
C Detailed acceptance–cost results
Table 3 compares draft-head designs through their acceptance–cost trade-off, using speedup, acceptance preservation, and relative LM-head cost across three datasets.
- Acceptance–cost comparison: Table 3 reports speedup relative to the full-vocabulary baseline, acceptance preservation ρτ, and relative LM-head cost ν.Values are averaged across MT-Bench, HumanEval, and GSM8K for Llama-3.1-8B at temperature 0 and batch size 1.
D Extended Results Tables
The extended tables report speedup and average acceptance length across models, temperatures, batch sizes, and evaluated configurations, with repeated-run aggregation and full-vocabulary reference comparisons.
- Reporting protocol: Speedup is measured relative to vanilla inference without a speculator under the same benchmark, temperature, and batch size, while results are averaged over five repeated runs.The tables report mean ± standard deviation, and Avg is the arithmetic mean across three benchmarks for each batch size.
- Llama3.1-8B-Instruct: Speedup tables cover Llama3.1-8B-Instruct at temperature 0 with batch sizes 1 and 64, and at temperature 1 with batch size 1.These tables are Table 4, Table 6, and Table 8.
- Llama3.1-8B-Instruct: Acceptance-length tables cover Llama3.1-8B-Instruct at temperature 0 with batch sizes 1 and 64, and at temperature 1 with batch size 1.These tables are Table 5, Table 7, and Table 9.