Graph-Structured Speculative Decoding (2024)

Abstract: Speculative decoding has emerged as a promising technique to accelerate the inference of Large Language Models (LLMs) by employing a small language model to draft a hypothesis sequence, which is then validated by the LLM. The effectiveness of this approach heavily relies on the balance between performance and efficiency of the draft model. In our research, we focus on enhancing the proportion of draft tokens that are accepted to the final output by generating multiple hypotheses instead of just one. This allows the LLM more options to choose from and select the longest sequence that meets its standards. Our analysis reveals that hypotheses produced by the draft model share many common token sequences, suggesting a potential for optimizing computation. Leveraging this observation, we introduce an innovative approach utilizing a directed acyclic graph (DAG) to manage the drafted hypotheses. This structure enables us to efficiently predict and merge recurring token sequences, vastly reducing the computational demands of the draft model. We term this approach Graph-structured Speculative Decoding (GSD). We apply GSD across a range of LLMs, including a 70-billion parameter LLaMA-2 model, and observe a remarkable speedup of 1.70 \times to 1.94 \times , significantly surpassing standard speculative decoding.

## introduction the impressive performance of large language models (llms) comes with an efficiency bottleneck that hinders their broader adoption @xcite ## preliminaries: sequence-structured speculative decoding in this section, we establish the notation and provide a foundational overview of sequencestructured speculative decoding (ssd). consider an input sequence at time step t, denoted by @xmath0, x 2 , ..., x t }, where each x i symbolizes the i-th token from the sequence. let m p be the target llm we want to accelerate, and let m q denote the draft model. the probabilities p(x t+1 |@xmath1, respectively. ssd leverages the draft model, m q , to propose a hypothesis comprising γ tokens, which we denote as @xmath2, xt+2 , ..., xt+γ }. the drafting of each token, xt + i, is modeled as follows: upon completion of the draft stage, the llm verifies the γ drafted tokens in a singular forward pass. the verification process, which compares predictions made by m p and m q to determine which tokens shall be accepted, can be conducted in both deterministic and non-deterministic ways. deterministic verification accepts drafted tokens only if the llm would generate the same. the nondeterministic way employs the sampling method used in previous studies @xcite . for the i-th token in the hypothesis, the acceptance probability is calculated as min(1, p(x t+i )/q(x t+i )). should the token xt+i face rejection, all subsequent tokens in the hypothesis are also discarded, the verification process comes to a halt, and m p regenerates the discarded token. this method ensures that the tokens that are ultimately accepted are representative of the output distribution characterized by m p . ## a step forward: tree-structured speculative decoding an intuitive idea for improving ssd is to draft multiple hypotheses instead of merely one. this is where tree-structured sd (tsd) comes into play. in each drafting step of ssd, the draft model predicts a single next token as described in equation 1. after γ steps, the drafted tokens compose a sequence {x t+1 , xt+2 , ..., xt+γ }. in contrast, tsd allows the draft model to consider k different alternatives for the next token at each drafting step. the resulting drafted tokens thus create a tree structure, with the root representing the context at the commencement of drafting, and each branch from the root depicting a different hypothesis. after γ drafting steps, the resulting token tree has a depth of γ and a maximum out-degree of k and can contain up to k γ+1 -1 k-1 nodes, representing as many as k γ unique hypotheses. let's denote the collection of all hypotheses as {h i } k γ @xmath3. tsd holds a significant advantage over ssd; by enabling the generation of a larger pool of hypotheses in a single drafting stage, it raises the chances of having longer sequences of tokens accepted by the llm. this boosts the acceptance rate of the sd process. fundamentally, tsd operates in a manner analogous to beam search, maintaining multiple potential hypotheses within its tree structure during the draft stage and then selecting the most promising one during the verification stage. ## graph-structured speculative decoding empirically, we observe that tsd often fails to surpass ssd, contrary to expectations. it appears that despite the utilization of pruning and tree attention, the cost of drafting multiple hypotheses still counterbalances the potential benefits that tsd offers. so we would like to ask: can we further reduce the quantity of drafted tokens to enhance tsd's efficiency and effectiveness? ## conclusion in this paper, we introduce graph-structured speculative decoding (gsd), a novel decoding strategy that utilizes a token graph to concurrently record a multitude of sequence hypotheses within a single draft stage. we propose a redundant node merging technique and two pruning strategies to constrain the size of the token graph without unduly compromising the diversity of hypotheses. our extensive experiments demonstrate that gsd significantly increases the acceptance rate of drafted tokens while not introducing much computation, achieving a noticeable acceleration in speed compared to previous speculative decoding methods.