PaperlayerView exact revision on arXiv

Source-linked AI summary

LogitSpec: Accelerating Retrieval-based Speculative Decoding via Next Next Token Speculation

Tianyu Liu, Qitan Lv, Hao Li, Xing Gao, Xiao Sun, Xiaoyan Sun

arXiv:2507.01449v3cs.CL

TL;DR

Retrieval-based speculative decoding can struggle to find matched, accurate draft tokens, while draft-model-based approaches carry deployment and overhead costs. LogitSpec uses the last logit to speculate a next-next token and retrieves with both tokens, achieving 1.8×–2.6× speedup across evaluation benchmarks.

  • Problem

    Retrieval-based speculative decoding often struggles to find matched, accurate references for draft tokens, while draft-model-based speculative decoding adds deployment difficulties and overhead.

    Source: §1 · p. 1

  • Method

    LogitSpec uses the last logit to speculate a next-next token, then retrieves a relevant reference using both the next and next-next tokens.

    Source: §1 · p. 2

  • Results

    1.8×–2.6× speedup across all evaluation benchmarks is reported for LogitSpec.

    Source: §6 · p. 9

  • Takeaways & Limitations

    LogitSpec needs no additional draft model and is plug-and-play for integration into existing LLM frameworks.

    Source: §6 · pp. 8–9

    Real-world inference acceleration is less competitive, and speedups may be lower with short prompts because retrieval uses the prompt.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Speculative decoding (SD), where a small draft model is employed to propose draft tokens in advance and then the target model validates them in parallel, has emerged as a promising technique for LLM inference acceleration. Many endeavors to improve SD are to eliminate the need for a draft model and generate draft tokens in a retrieval-based manner in order to further alleviate the drafting overhead and significantly reduce the difficulty in deployment and applications. However, retrieval-based SD relies on a matching paradigm to retrieval the most relevant reference as the draft tokens, where these methods often fail to find matched and accurate draft tokens. To address this challenge, we propose LogitSpec to effectively expand the retrieval range and find the most relevant reference as drafts. Our LogitSpec is motivated by the observation that the logit of the last token can not only predict the next token, but also speculate the next next token. Specifically, LogitSpec generates draft tokens in two steps: (1) utilizing the last logit to speculate the next next token; (2) retrieving relevant reference for both the next token and the next next token. LogitSpec is training-free and plug-and-play, which can be easily integrated into existing LLM inference frameworks. Extensive experiments on a wide range of text generation benchmarks demonstrate that LogitSpec can achieve up to 2.61 $\times$ speedup and 3.28 mean accepted tokens per decoding step. Our code is available at https://github.com/smart-lty/LogitSpec.

1 Introduction

LogitSpec uses the last logit to speculate a next-next token and adds it as a retrieval cue to improve retrieval-based speculative decoding. The paper reports speedups up to 2.61× and 3.28 mean accepted tokens per decoding step without an extra draft model.

1 Introduction

  • Using a small draft model for speculative decoding adds deployment and implementation difficulties, memory overhead, and training cost when no draft model is available.Speculative decoding proposes draft tokens for parallel target-model validation while preserving the output distribution.Source: §1 · p. 1
  • Retrieval-based speculative decoding can fail to find matched, accurate draft tokens; PLD has no matched tokens in more than 30% of decoding steps.The introduction illustrates a failure where retrieval from “the” selects “triangle” instead of “area.”Source: §1 · p. 2, Figure 1 · p. 2
  • LogitSpec first speculates the next-next token from the last logit, then retrieves references using both the next and next-next tokens.This expands the retrieval range and can help filter redundant references or search when no relevant reference exists.Source: §1 · p. 2
  • The last logit can speculate the next-next token with relatively high accuracy, and this property is described as robust across tasks.The paper presents this observation as the basis for LogitSpec.Source: §1 · p. 2
  • Up to 2.61× speedup and 3.28 mean accepted tokens per decoding step are reported without an extra draft model.Experiments span various text-generation benchmarks.Source: §1 · p. 2

2 Related Work

Prior speculative decoding methods either rely on specialized draft models or eliminate them, including through partial-layer and retrieval-based approaches. LogitSpec adds a speculated next-next token as a retrieval anchor, a training-free mechanism the paper describes as complementary to methods with trained extra heads.

2 Related Work

  • Draft-model-based methods improve draft quality through specialized models, but require extra parameters or extensive training and can add deployment, memory, and retraining burdens.The paper discusses these costs especially for long-context settings and when the target model changes.Source: §2 · p. 3
  • Draft-model-free approaches include partial-layer drafting and retrieval; retrieval-based methods dominate this line, but context retrieval may miss matches and external-database retrieval may return inaccurate drafts.Partial-layer methods are limited by the number of layers they can skip.Source: §2 · p. 3
  • LogitSpec is described as the first method to use a speculated next-next token as an additional retrieval anchor, complementing multi-token prediction methods that require trained extra heads.The mechanism is training-free.Source: §2 · p. 3

3 Background

Speculative decoding drafts multiple tokens and verifies them together, while retrieval-based methods use matched n-grams to propose drafts without a separate draft model. Longer query matches are more precise but less likely to be found.

3 Background

  • A speculative decoding step drafts γ tokens and has the target model verify them in parallel; rejected drafts are resampled, while accepted drafts can yield up to γ + 1 tokens per step.This preserves the output distribution while enabling multiple tokens to be generated per target-model forward pass.Source: §3 · p. 3
  • Retrieval-based speculative decoding stores corpus n-grams as references and returns tokens following references that match the query m-gram.The matched references supply draft tokens without an additional draft model.Source: §3 · pp. 3–4
  • In over 50% of decoding steps, the next-next token appears among the top 60 entries of the last logit across model sizes and architectures.This observation motivates using the last logit to guide retrieval.Source: Figure 2 · p. 4
  • A larger query length m makes matches more precise but lowers the probability of finding a match.A match requires all m query tokens to agree with the corresponding reference tokens.Source: §3 · p. 4

4 Method

LogitSpec uses the last logit to speculate the next-next token and guide retrieval for both the next and next-next tokens. Its motivating experiments find this signal across models and tasks, and its drafting and verification stages organize retrieved sequences for parallel checking.

4 Method

  • LogitSpec is presented as a retrieval-based speculative-decoding framework that uses the last logit to guide retrieval and improve retrieval performance.The method section frames its motivating experiment, drafting procedure, and verification preparation as parts of this framework.Source: §4 · p. 4

4.1 Motivated Observation

  • In over 50% of decoding steps, the next-next token appears in the last logit's top 60 entries across models and tasks; for Llama 3.1 8B, the figure is 64%.The motivating experiment covers six models, including Vicuna, Llama 3.1, and Qwen, ranging from 0.5B to 33B parameters.Source: §4.1 · p. 4
  • Last-logit decoding achieves mean accepted tokens per decoding step (MAT) above 1.5 and shows greater task robustness than PLD and REST.The experiment uses Vicuna 7B and supplies top-60 last-logit entries as draft tokens, each a guess at the next-next token.Source: §4.1 · p. 4

4.2 LogitSpec Drafting

  • LogitSpec samples the next token, uses top-k entries of the last logit other than that token to speculate the next-next token, then retrieves references for both.It first tries retrieval with m = 3, decreasing m when no match is found; the retrieval model uses the user prompt and decoded tokens, making it query-independent.Source: §4.2 · p. 5
  • Retrieved draft sequences are organized into a draft tree, with pruning that preserves more retrieved tokens for higher-ranked next-next-token speculations.The pruning heuristic preserves four tokens below rank 8, three below rank 32, and otherwise only the speculated token, until capacity K is exceeded.Source: §4.2 · p. 5, §4.3 · p. 5

4.3 LogitSpec Verification

  • For parallel verification, LogitSpec prepares an attention mask that makes each draft-token sequence invisible to the others.The next token remains visible to all sequences; the mask uses causal masks sized to each sequence.Source: §4.3 · p. 5

5 Experiments

LogitSpec is evaluated across text-generation benchmarks using mean accepted tokens per decoding step and overall speedup. It outperforms retrieval-based baselines, and ablations and case studies examine its components, retrieval performance, and runtime overhead.

5.1 Experimental Setup

  • The evaluation spans Spec-Bench and benchmarks for code, mathematics, summarization, and long-context generation, using MAT and overall speedup as metrics.The setup follows standard inference settings with temperature 0 and batch size 1.Source: §5.1 · pp. 5–6

5.2 Main Results

  • 2.61× speedup on CNN/DM with Vicuna 7B and 2.47× on HumanEval with Vicuna 13B; LogitSpec outperforms PLD across benchmarks and model sizes.It also remains strong with limited context repetition, where other retrieval-based methods struggle; on CNN/DM, REST achieves less than 1.5× speedup.Source: Table 1 · p. 7, §5.2 · p. 7

5.3 Ablation Study

  • Ablating either last-logit decoding or retrieval degrades LogitSpec; removing retrieval lowers MAT by 0.87 on Spec-Bench with Vicuna 7B and 1.71 on CNN/DM, while last-logit removal causes more consistent MAT degradation across tasks.Varying the pruning strategy yields only minor differences.Source: §5.3 · p. 7

5.4 Case Study

  • LogitSpec retrieves a matched reference in most decoding steps, whereas PLD has no matched token in more than 30% of decoding steps.The retrieval-success-rate comparison covers Spec-Bench and HumanEval with Vicuna.Source: Table 3 · p. 8, §5.4 · p. 8
  • Model forward occupies most wall-clock time, while retrieval overhead is 1.17% of the whole decoding process.The retrieval process is independent and can be further alleviated by parallel techniques.Source: Figure 5 · p. 8, §5.4 · p. 8

6 Conclusion

The last token’s logit can predict the next-next token without fine-tuning, and LogitSpec uses that prediction to expand retrieval. The plug-and-play method improves retrieval performance and achieves 1.8×–2.6× speedup across evaluation benchmarks.

6 Conclusion

  • LogitSpec uses the last logit’s next-next-token prediction to expand retrieval and find relevant references as draft tokens, without an additional draft model.The last logit can predict the next-next token without fine-tuning.Source: §6 · p. 8
  • LogitSpec is fully plug-and-play and can be integrated into existing LLM frameworks.The paper describes it as easy to implement and integrate.Source: §6 · p. 9
  • 1.8×–2.6× speedup across all evaluation benchmarks accompanies improved retrieval performance.Source: §6 · p. 9

Limitations

The authors identify limits in real-world acceleration, prompt-based retrieval, and the conditions used to estimate speedups.

Limitations

  • LogitSpec’s real-world inference acceleration is less competitive, and the authors plan to integrate it with draft-model-based methods.Retrieval from prompt context may yield lower speedup on short prompts; benefits also fall on out-of-domain corpora.Source: Limitations · p. 9
  • Speedup estimates come from offline decoding with curated prompts and NVIDIA A100 GPUs, and may vary across workloads or hardware.Translation tasks with scarce context n-grams still achieve ∼1.4× speedup, while LongBench with longer contexts reaches 2.01×.Source: Limitations · p. 9

Ethical Considerations

LogitSpec inherits the ethical properties and risks of its underlying models. The related-work discussion surveys draft-model-based and draft-model-free approaches, including retrieval methods.

Ethical Considerations

  • LogitSpec inherits the biases and potential safety risks of its underlying models; it neither adds harmful-content capabilities nor mitigates existing biases.The authors recommend standard safety guardrails and alignment techniques for deployed target models.Source: Ethical Considerations · p. 9

A More Discussions to Related Work

  • Draft-model-based work addresses draft quality, verification, and adaptive draft length through approaches including HASS, CORAL, Judge Decoding, Gumiho, and LongSpec.The cited methods target training-inference inconsistency, representation alignment, acceptance of refused tokens, hybrid drafting, and long-context decoding, respectively.Source: A More Discussions to Related Work · p. 13
  • Draft-model-free work includes layer-sparsity and retrieval approaches; Token Recycling uses dropped tokens and an adjacency matrix, while SAM Decoding retrieves from text sources using a suffix automaton.The authors note Token Recycling is query-dependent and may be limited with complex, dynamic user inputs.Source: A More Discussions to Related Work · p. 13

B Details of Motivation Experiments

The appendix examines last-logit prediction, implementation choices, baselines, and experimental settings. It also reports that retrieval accounts for a small share of measured runtime after optimization.

B Details of Motivation Experiments

  • Motivation experiments assess last-logit prediction across six models and test robustness on a small Spec-Bench subset; the authors offer intuition but leave formal theory to future work.They report that top-ranked last-logit tokens correlate with likely next tokens, particularly in structured text, without fine-tuning.Source: B Details of Motivation Experiments · pp. 13–14
  • The ground-truth next-next-next token appears among the last logit’s top 60 tokens with ≥40% probability, but extending speculation was rejected due to tree growth.Speculating top 60 for both next-next and next-next-next tokens would create at least 3,600 draft tokens; capacity beyond 64–128 yields diminishing returns as verification overhead dominates.Source: B Details of Motivation Experiments · p. 13

C Pseudo Code to Prepare Attention Inputs

  • LogitSpec organizes retrieved draft sequences as a tree and prepares a local causal mask for each subsequence.The pseudo-code prepares attention inputs for the retrieved candidates.Source: C Pseudo Code to Prepare Attention… · p. 14

D.1 Baselines

  • The comparison includes plug-and-play retrieval baselines Lookahead Decoding, REST, and PLD, plus vanilla speculative sampling with Vicuna-68M.The main baseline comparisons use Vicuna and Llama-2 models.Source: §D.1 · p. 15

D.2 Hardware and Software Configurations

  • Experiments use float16, temperature 0, batch size 1, and a maximum generation length of 1024; 70B models run on two A100 80G GPUs.Models from 7B to 33B use one A100 80G GPU.Source: §D.2 · p. 15

D.3 Retrieval Process

  • Hash-based retrieval reduces per-step work to O(n + k) after amortized O(L) prompt preprocessing, instead of scanning the prompt for each candidate.Each lookup uses an m-gram query and returns a continuation span; the reported total per-step cost is O(k + kn).Source: §D.3 · p. 15
  • 1.17% of total wall-clock time is attributed to retrieval in Figure 5, supporting the effectiveness of the authors’ optimization.The hash table’s O(mL) memory cost is described as negligible compared with the KV cache.Source: §D.3 · p. 15

D.4 Evaluation Instructions

  • Evaluation uses task-specific instructions across Spec-Bench, HumanEval, CNN/DM, and GSM8K.Examples include code-completion instructions for HumanEval and a “Summarize:” prefix for CNN/DM.Source: §D.4 · pp. 15–16

D.5 Dataset Configurations

The appendix specifies evaluation and model configurations, ethics practices, and an additional comparison with draft-model-based methods.

D.5 Dataset Configurations

  • Spec-Bench and HumanEval use their full data, while CNN/DM and GSM8K use randomly sampled sets of 1,000 questions.Maximum generation length is 1,024 across tasks.Source: §D.5 · p. 16

D.6 Model Configurations

  • Models use float16, tokenizer-matched EOS tokens, and continuous GPU memory allocation for KV caches to alleviate rollback overhead.The continuous allocation follows Medusa.Source: §D.6 · p. 16

Ethics Statement

  • The study uses publicly available datasets under their licenses, involves no human subjects or personally identifiable information, and advises against high-stakes deployment without further safety assessments.The authors state that the work adheres to the ACL Code of Ethics.Source: Ethics Statement · p. 16

E.1 Comparison with Draft-Model-Based Methods

  • On Spec-Bench with Vicuna-7B, LogitSpec’s speedup is comparable to EAGLE’s and above Medusa’s, while requiring no training or extra parameters.The table caption notes that EAGLE and Medusa require training extra heads.Source: Table 5 · p. 16, §E.1 · p. 16

E.2 More Results on Different LLM Backbones

LogitSpec outperforms retrieval-based baselines on the reported comparisons, with stronger speedup on AIME, a low-repetition reasoning dataset. Additional Llama-2 results show best acceleration across the tested baselines.

Comparison with State-of-the-Art Retrieval-Based Methods

  • Across AIME, LongBench, and MATH on Llama-3.1-8B-Instruct, LogitSpec achieves 7.5% higher average speedup, 2.58× versus Token Recycling’s 2.40× and SAM’s 2.34×.Source: §E.2 · p. 17
  • LogitSpec is query-independent, unlike Token Recycling’s use of previously rejected drafts and cross-query information, and is described as more robust to diverse query patterns.The passage contrasts this with Token Recycling’s query-dependent adjacency-matrix approach.Source: §E.2 · p. 17
  • Speculating next-next tokens expands retrieval beyond observed n-grams, addressing SAM Decoding’s low match rates when context repetition is limited.SAM uses suffix automata for n-gram matching from a static corpus.Source: §E.2 · p. 17
  • On AIME, LogitSpec achieves 3.09× speedup versus Token Recycling’s 2.82× on a challenging reasoning dataset with low repetition.The passage presents this as demonstrating the value of predictive guidance when retrieval alone struggles.Source: §E.2 · p. 17

Additional LLM backbones

  • Across Spec-Bench, CNN/DM, GSM8K, and HumanEval, LogitSpec consistently achieves the best acceleration among the reported Llama-2 baselines.The results cover Llama-2 chat models; Tables 8 and 9 report the evaluations.Source: §E.2 · p. 17

E.3 Per-Task Spec-Bench Breakdown on Vicuna

Across additional tasks and settings, LogitSpec retains strong speedups, including under stochastic decoding and long contexts. The appendix also examines cross-lingual predictive coverage, pruning and draft-tree capacity, a worked example, and production-framework integration.

E.3 Per-Task Spec-Bench Breakdown on Vicuna; E.4 More Results on Different Benchmarks

  • Without external databases, LogitSpec reaches 1.38×–1.43× speedup on Translation versus REST’s 1.17×–1.31×; broader evaluations report around 2× average acceleration.On MATH it reports 3.32 MAT and 2.78× speedup, and on AIME25 up to 3.76 MAT and 3.33× speedup.Source: §E.3 · p. 17, §E.4 · p. 17

E.5 Sampling with Non-Zero Temperatures

  • At T=0.8, LogitSpec still achieves speedups above 1.8×; higher temperatures reduce acceptance rates for speculative decoding methods generally.Speculative candidates are selected from raw logits before temperature scaling, so the candidate set is unchanged across temperatures.Source: §E.5 · pp. 17–18

E.6 Cross-Lingual Next-Next Token Coverage

  • At k=64, over 50% of ground-truth next-next tokens appear among the last logit’s top-k predictions in each of five MMMLU languages.The languages tested are English, Chinese, Japanese, German, and French.Source: §E.6 · p. 20, Table 14 · p. 21

E.8 Ablation Study on Different Pruning Strategies

  • Across pruning strategies, speedup remains stable at approximately 1.9×, while pruning controls draft-tree size to balance breadth and cost.The authors describe pruning as auxiliary to next-next token speculation, the core mechanism.Source: §E.8 · p. 21

E.7 More Results on Long-Context Scenarios

  • 3.09 MAT and 2.01 overall speedup are reported for LongBench long-context scenarios.The hash-table retrieval overhead is around 1.17% of total decoding time; next-next token speculation helps disambiguate retrieval among repetitive phrases.Source: §E.7 · pp. 20–21, Figure 7 · p. 22

E.9 Ablation Study on Draft Tree Capacity

  • MAT rises from 2.03 to 2.30 as K increases, but speedup peaks at 1.92× with K=64 and then declines.The authors attribute the decline to verification overhead from an overly large draft tree.Source: §E.9 · p. 22, Table 16 · p. 23

E.10 Real-World Example

  • In the worked example, last-logit guidance allows acceptance of three draft tokens, while standard retrieval retrieves only “A pencil costs.”The accepted sequence is “the combined cost of one pen”.Source: §E.10 · pp. 22–23

F Discussion: Integration with Production Frameworks

  • vLLM integration has no fundamental algorithmic obstacles, but engineering integration remains future work; LogitSpec’s draft tree fits shared-prefix batching.The discussion notes that vLLM has native tree-attention support and that retrieval and tree construction can run on the CPU alongside GPU inference.Source: F Discussion: Integration with… · p. 23
Loading 2507.01449v3…