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

TL;DR

LogitSpec addresses the challenge of accelerating retrieval based speculative decoding by leveraging the last token logit's ability to speculate the next next token. It introduces a two step drafting process that first speculation the next next token and then retrieves references for both the next and next next tokens, organized into a draft tree with a lightweight tree attention scheme for parallel verification. Across a wide range of benchmarks and backbones, LogitSpec achieves up to $2.61\times$ speedup and improves MAT to $3.28$, all without a dedicated draft model. The approach is plug and play, robust to long contexts and language diversity, and offers a practical path for deploying faster LLM inference while maintaining lossless output quality.

Abstract

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.

Figures (7)

• Figure 1: Illustration of vanilla retrieval-based SD method and our LogitSpec. (a) Retrieval-based SD retrieves the wrong token "triangle" according to the next token "the". (b) LogitSpec first speculates the next next token "area", and then retrieves the right relevant reference "area of" according to "the area". This simple example illustrates how LogitSpec utilizes the last logit to speculate the next next token and improves the retrieval accuracy.
• Figure 2: Motivated observation I: the last logit can speculate the next next token with high accuracy. In over 50% decoding steps, the next next token can be found in the top-60 entries within the last logit across model sizes and architectures.
• Figure 3: Motivated observation II: compared with other retrieval-based methods, the prediction of the last logit demonstrates robustness to downstream tasks, motivating us to utilize it to guide the retrieval process.
• Figure 4: An overview of LogitSpec. At each decoding step, LogitSpec first utilizes the top-$k$ entries of the last logit as the speculation for the next next token. Then, LogitSpec retrieves relevant references for both the next token and the next next token. Finally, LogitSpec organizes the draft tokens into a draft tree and prepares a tree attention for parallel verification.
• Figure 5: Running time breakdown of the whole decoding process on Spec-Bench with Vicuna 7B.