Multi-Token Joint Speculative Decoding for Accelerating Large Language Model Inference

Large Language Models (LLMs) have exhibited remarkable performance across real-world tasks spanning text and image domains [4, 7, 33, 34]. However, the substantial parameter size of these models leads to high computational cost during inference [27]. For example, ChatGPT is estimated to have an energy demand of 564 MWh per day [8]. Compounding this issue, the autoregressive generation of $K$ tokens necessitates calling the model $K$ times, involving the repeated loading of weight matrices and intermediate results from GPU global memory to computing units. As a result, LLM inference is often hampered by constraints related to memory bandwidth and communication [20], leading to diminished hardware utilization, heightened inference latency, and elevated energy cost.

To tackle this challenge, researchers have delved into non-autoregressive decoding approaches. Earlier methods [11, 12, 13] aimed at reducing inference latency by concurrently generating multiple tokens. But these methods usually require task-dependent techniques and information to achieve comparable performance to that of autoregressive decoding [17, 36]. More recently, speculative decoding has emerged [20, 5, 17, 32], exploiting the observation that, even when a small model is ten times smaller than a large model, most of the small model’s prediction aligns well with that of a large model. So it leverages a smaller auxiliary model to draft a few future tokens autoregressively, which are subsequently validated in parallel by the larger model. As the smaller model operates significantly faster and parallel token verification incurs a similar time cost as generating a single token, speculative decoding achieves an overall speed-up of $1$-$2\times$. Meanwhile, it always has an identical sampling distribution as greedy decoding.

However, there are some limitations of the existing studies on speculative decoding. First, they aim to have an identical sampling distribution as greedy sampling. But the goal of decoding algorithms is to generate sequences with optimal perplexity, which is defined as the exponentiated average negative log-likelihood of all tokens. An algorithm that does not follow the same sampling distribution as greedy decoding, but yields output with better perplexity could be more useful in the real world. Second, in these algorithms, if a draft token is rejected, all the draft tokens after it are discarded. But the overall draft sequences may have good quality despite there being a few low-quality tokens. Rejecting these sequences will shorten the average acceptance length for each iteration, lowering the decoding efficiency. Third, existing research does not investigate the impact of speculative decoding on inference energy consumption. While speculative decoding increases the number of FLOPs during inference, it simultaneously reduces the overall inference time, communication, and memory operations, which are crucial factors in determining energy consumption [1, 6]. Thus, it remains unclear if speculative decoding increases or decreases the energy cost of inference.

So we aim to design a new algorithm that is both more efficient and more effective than greedy decoding and vanilla speculative decoding. We first introduce multi-token joint greedy decoding (MJGD), which generates multiple tokens at each step greedily based on their joint likelihood. We empirically show that it has better overall perplexity than greedy decoding does. However, the computational cost of MJGD is infeasible in practice. So we further propose multi-token joint speculative decoding (MJSD), which approximates and accelerates MJGD from two aspects: (1) it approximates the joint distribution of the large model with that of a small model by generating draft tokens with the small model and validating them with the large model; (2) it uses beam decoding to accelerate the sequence generation from the joint distribution of the small model. Compared with vanilla speculative decoding, MJSD has two advantages: (1) it is more effective, as it approximates MJGD, which has better output perplexity than greedy decoding does; (2) it is more efficient because it uses joint likelihood as the verification criterion, which allows it to accept the longest prefix sub-sequence of the draft tokens with valid perplexity. In addition, we analyze the energy consumption of LLM inference. We give theoretical and empirical evidence that, despite that MJSD and other speculative decoding algorithms increase the number of FLOPs needed during LLM inference, they reduce the overall energy consumption by reducing the overhead induced by accessing GPU global memory.

We evaluate our method on text generation tasks with Llama-2 [34] and OPT models [38]. Experiment results show that our approach enhances perplexity by $21.2\%$ than the baselines with identical sampling distributions as greedy decoding. In addition, MJSD achieves $2.21\times$ speed-up and $2.84\times$ smaller energy consumption than greedy decoding, and $1.49\times$ speed-up and $1.62\times$ smaller energy consumption than vanilla speculative decoding. Our code is open-sourced¹¹1https://anonymous.4open.science/r/LLMSpeculativeSampling-EE52.

The perplexity is a direct measurement of the effectiveness of a decoding algorithm. In practice, when the model is well-trained, a better perplexity usually leads to better downstream effectiveness. Although it is possible that improving perplexity does not improve the performance on downstream tasks, the problem is due to the model itself instead of the decoding algorithm. Therefore, the goal of this work is to design an algorithm that yields better output perplexity and better efficiency than greedy decoding and vanilla speculative decoding.

In this section, we first introduce multi-token joint greedy decoding (MJGD) to give a motivating example for why generating multiple tokens based on their joint likelihood can lead to better perplexity. Then we introduce multi-token joint speculative decoding (MJSD), which approximates and accelerates MJGD by exploiting a small model.

Previous studies [20, 5, 17, 32] only focus on the speed of speculative decoding. However, an equally important consideration is energy consumption. To our knowledge, there is no existing work evaluating the impact of speculative decoding on inference energy consumption. Although MJSD and speculative decoding increase the number of FLOPs due to the involvement of an additional small model and the rollback operation, they concurrently reduce the inference time and memory operations, which are key factors of GPU energy consumption [1, 6]. Consequently, it poses an open question regarding whether speculative decoding increases or decreases overall energy consumption.

To understand the net effect of speculative decoding, we decompose the total energy consumption into two parts following [1]:

where $PW_{flop},PW_{mem}$ denote the power (energy/second) of FLOPs and memory operations, and $T_{flop},T_{mem}$ denote the time spent on these operations. When input batch size increases, $PW_{flop}$ increases until it reaches the power of maximum FLOPs, denoted as $PW^{*}_{flop}$. Meanwhile, $PW_{mem}$ is irrelevant to the input batch size, as it only depends on the memory hardware.

To determine the relative magnitude relationship between $PW_{flop}$ and $PW_{mem}$, we first point out the fact that GPU memory operations in LLM inference are dominated by accessing off-chip global memory, which consumes about $100\times$ of energy compared to accessing on-chip shared memory [16]. It is because each multiprocessor on GPU usually has 64KB of on-chip memory shared by multiple threads, while storing a single layer of LLM, say T5-11b [25], requires about 1GB memory. Moreover, Allen and Ge showed that doing sequential read from off-chip memory consumes $20$-$30\%$ more power than running maximum FLOPs [1]. So we have $PW_{mem}>PW^{*}_{flop}\geq PW_{flop}$. Notice that $PW^{*}_{flop}=PW_{flop}$ only if the batch size reaches the maximum parallelization capacity of GPUs. During greedy decoding and speculative decoding, the batch size is usually small [20]. So most of the computing power is not utilized [20], which means $PW_{mem}\gg PW_{flop}$.

In addition, previous studies have shown that during LLM inference $T_{mem}\gg T_{flop}$ [20]. Therefore, the energy induced by memory operations, i.e., $PW_{mem}T_{mem}$ dominates $E_{total}$. Since speculative decoding reduces $T_{mem}$ by reducing the number of runs of the large model, it should reduce the inference energy consumption to a similar extent as it reduces time consumption.

To validate our hypothesis, we conducted an experiment to evaluate how batch size influences energy consumption during inference. We ran OPT-13b models on 2 Nvidia A10 GPUs with 24GB memory. Fixing the total number of runs of the large model while varying the input batch size $b\in\{1,2,4,8,16\}$ for each run, we measured time and energy cost. The details of energy measurement are illustrated in the Appendix B. Table 1 shows the results. As batch size doubles, although the number of FLOPs doubles, the energy consumption per run increases slightly. This observation demonstrates that $PW_{mem}T_{mem}$ dominates $E_{total}$. In the next section, we will show more direct evidence for the energy efficiency of MJSD and other speculative decoding algorithms.

Datasets and Models.We use three public datasets for evaluation: (1) ChatGPT-Prompt [26], (2) ChatAlpaca [3], and (3) CNN Dailymail [30]. The three datasets represent three common tasks for LLMs: instruction, multi-turn chat, and summarization. Table 4 in the Appendix shows more details of the datasets. We use two pairs of small and large models in our experiments: OPT-125m and OPT-13B [38], and Llama-68M [24] and Llama-2-13B [34].

Baselines. For each pair of small and large models, we compare our method with greedy decoding (greedy), and three speculative decoding methods: vanilla speculative decoding (speculative) [19, 5], Spectr [32], and BiLD [17]. Our implementation of MJSD and all the baselines are based on a public implementation of speculative decoding [2]. We also tried to compare with SpecInfer [24] using their released implementation. But it runs out of memory on our machine. For each method, we let it generate at most 128 tokens for each input and run it for $1,000$ seconds. We open-sourced our code for reproduction.

More details of the hyper-parameters, warping operations, and machine configurations of the experiments can be found in the Appendix B, C, and D .

Efficient Decoding Inference. There are extensive studies on improving large model inference efficiency. Well-known methods include model quantization [9, 22], model pruning [10, 28], and model distillation [15]. Despite achieving significant speed-ups, a common drawback of these methods is that they have to sacrifice the model’s effectiveness.

A direction closer to our work is non-autoregressive decoding. It is first proposed by [12] to generate multiple tokens in parallel. That is, the model simultaneously predicts $p(x_{t+k}|x_{1:t})$ ($k=1,2,\ldots$). Subsequent studies further improved the performance of parallel decoding by incorporating additional information [35, 31, 21] or using additional iterations to refine predictions [11, 19, 13]. However, these works require continuous training of the model and usually either compromise the model effectiveness or require task-dependent techniques to achieve comparable performance [17].

Speculative Decoding. Speculative decoding was recently proposed in [20, 5] as a way to accelerate LLM inference. Spectr [32] enhances speculative decoding by letting the small model generate multiple i.i.d. draft sequences. While speculative decoding and Spectr use the large model to verify all the tokens drafted by the small model, BiLD [17] only calls the large model when the probability output by the small model is below a pre-defined threshold $\tau_{1}$. The large model rejects a token if its negative log-likelihood is larger than threshold $\tau_{2}$. SpecInfer [24] uses one or multiple small models to generate a draft token tree to increase the average acceptance length for each iteration. All these methods can be perceived as exact or approximate versions of sampling tokens from the conditional distribution $p(x_{t}|x_{<t})$. Therefore, their output perplexity is bounded by greedy decoding.

An orthogonal direction to improve speculative decoding is to improve the effectiveness of the small draft model. It is obvious that if more draft tokens are accepted, the overall inference speed will increase. BiLD [17] uses a model prediction alignment technique to better train the small model. Liu et al. [23] propose online speculative decoding to continually update the draft model based on observed input data. Instead, Rest [14] uses a retrieval model to produce draft tokens. The improvement of small models can benefit all speculative decoding algorithms, including our method.

We introduce multi-token joint speculative decoding that significantly enhances output quality with better time and energy efficiency. A distinctive aspect of our work is the exploration of speculative decoding’s impact on inference energy consumption, an often neglected area in existing studies. Experiment results demonstrate our method achieves significant energy reduction. This research contributes not only a novel decoding approach but also valuable insights for optimizing LLM deployment in real-world applications where considerations of both quality and efficiency are crucial.

Limitations. This paper mainly considers improving the output perplexity of decoding algorithms. Under the assumption that the model is well-trained for downstream tasks, improving output perplexity usually leads to better downstream effectiveness. But if the model is not well-trained, improving the perplexity may not necessarily improve the downstream effectiveness. However, it is mainly the problem of the model itself. In real-world applications, it is reasonable to assume the model is properly trained for downstream tasks.

Impact Statements. The goal of this work is to advance the field of Large Language Model (LLM) Inference, which has received tremendous attention from both academia and industry. However, LLMs also have received many critiques, including their extremely large carbon footage emission during both training and inference. Our work pays special attention to their energy consumption during inference to provide high-quality and fast-inference LLMs with reduced energy consumption, which has become a serious concern with the rapid increase of LLM-related workloads in the past few years.