Multi-Token Joint Speculative Decoding for Accelerating Large Language Model Inference
Abstract
Transformer-based Large language models (LLMs) have demonstrated their power in various tasks, but their inference incurs significant time and energy costs. To accelerate LLM inference, speculative decoding uses a smaller model to propose one sequence of tokens, which are subsequently validated in batch by the target large model. Compared with autoregressive decoding, speculative decoding generates the same number of tokens with fewer runs of the large model, hence accelerating the overall inference by -. However, greedy decoding is not the optimal decoding algorithm in terms of output perplexity, which is a direct measurement of the effectiveness of a decoding algorithm. An algorithm that has better output perplexity and even better efficiency than speculative decoding can be more useful in practice. To achieve this seemingly contradictory goal, we first introduce multi-token joint greedy decoding (MJGD), which greedily generates multiple tokens at each step based on their joint perplexity. We show that it leads to better perplexity for the whole output. But the computation cost of MJGD is infeasible in practice. So we further propose multi-token joint speculative decoding (MJSD), which approximates and accelerates the MJGD from two aspects: it approximates the joint distribution of the large model with that of a small model, and uses a verification step to guarantee the accuracy of approximation; then it uses beam decoding to accelerate the sequence generation from the joint distribution. Compared with vanilla speculative decoding, MJSD has two advantages: (1) it is an approximation of MJGD, thus achieving better output perplexity; (2) verification with joint likelihood allows it to accept the longest prefix sub-sequence of the draft tokens with valid perplexity, leading to better efficiency. In addition, we analyze energy consumption during LLM inference and provide theoretical and empirical evidence that, surprisingly, MJSD reduces energy consumption even when the number of operations increases. Experiment results show that our approach enhances perplexity by than greedy decoding. Moreover, MJSD achieves speed-up and less energy consumption than greedy decoding, and speed-up and less energy consumption than vanilla speculative decoding.
1 Introduction
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 tokens necessitates calling the model 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 -. 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 than the baselines with identical sampling distributions as greedy decoding. In addition, MJSD achieves speed-up and smaller energy consumption than greedy decoding, and speed-up and smaller energy consumption than vanilla speculative decoding. Our code is open-sourced111https://anonymous.4open.science/r/LLMSpeculativeSampling-EE52.
2 Preliminaries
2.1 Decodings of LLMs
Let denote the distribution defined by LLM model . Given an input prefix , the optimal decoding algorithm is to generate a sequence of tokens with maximum likelihood . The likelihood of the sequence is directly linked to perplexity of the sequence, which is the exponentiated average negative log-likelihood of all tokens. Based on autoregressive decomposition , the perplexity is defined as:
(1) |
Now we introduce commonly used decoding approaches.
Greedy Decoding.
Greedy decoding samples the next single token based on . According to the autoregressive decomposition, greedy sampling leads to the desired joint distribution for the sequence. In practice, however, the objective of decoding algorithms is to generate sequences with high likelihood. Instead of sampling from the joint distribution, a warping operation usually is applied to boost the high probability region, which transforms the original distribution to . A very popular warping operation is , which only selects the sample with the highest probability. Other possible warping operations include top-k. For most of the warping operations, autoregressive decomposition no longer holds after the warping. For example, .
Beam Decoding.
Beam decoding aims to do a better job than greedy decoding. For each position (), it maintains candidate sequences, which are also called beams. Still take warping operator as an example and assume we have already kept the sequences at position with highest likelihood, tokens are then sampled from for each sequence and the top sequences with the highest likelihood will be kept. It is widely recognized that beam decoding is a closer approximation to than greedy decoding [18], with few exceptions under special construction. The computation for each beam can be parallelized, thus its run time is about the same as greedy decoding. The fact that beam decoding is more effective than greedy decoding also suggests that the output perplexity of greedy decoding still has room for improvement.
2.2 Vanilla Speculative Decoding
Besides effectiveness, speculative decoding is proposed by [20, 5] to accelerate the inference of LLMs. It utilizes a small model to generate the next tokens and then uses the large model to verify the drafted tokens in parallel, which is summarized below:
-
1.
Let be the input context, the small model samples draft tokens using greedy decoding based on the warped predicted conditional probability for , where and is the small model’s output distribution.
-
2.
The large model verifies the draft tokens in parallel by computing the conditional probability for .
-
3.
Each draft token is accepted with probability . The draft tokens before the first rejected token are kept as the decoding output. An additional token is sampled from a residual distribution as a correction to the first rejected token. Then the accepted tokens and the resampled token are appended to the context as the input to the next iteration.
-
4.
Repeat step 1-3 until reaching the stopping criteria, e.g., reaching the length limit..
Because the large model verifies tokens in parallel with one run, the time cost is smaller than calling it times. Meanwhile, although the small model still runs in an autoregressive way, its inference speed is much faster than the large model. So speculative decoding can accelerate the inference process of LLMs. Additionally, it is proven that each token generated by speculative sampling follows the identical sampling distribution as greedy decoding [20]. However, since , the effectiveness of speculative decoding also has room for improvement. Therefore, instead of having an identical sampling distributon as greedy decoding, our algorithm aims to achieve a better approximation to .
3 Methodology
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.
3.1 Multi-Token Joint Greedy Decoding
Let be the large target model, the goal is to generate a sequence of tokens with maximum likelihood . We now introduce a new decoding algorithm to improve greedy decoding in terms of perplexity.
Definition 3.1.
Multi-Token Joint Greedy Decoding. Let be the large target model with distribution . Different from single-token greedy decoding, multi-token joint greedy decoding (MJGD) generates the next tokens at step based on their joint conditional probability 222In the paper, we omit when there is no ambiguity., where is an integer no less than 1 and , i.e., the total tokens generated in the previous steps.
![Refer to caption](extracted/5728708/figures/k-tokenGD.png)
![Refer to caption](extracted/5728708/figures/k-tokenGD-rouge.png)
Greedy decoding is a special case of MJGD where , . When , MJGD generates the optimal sequence by returning the sequence with the highest likelihood. So intuitively, output perplexity should improve as increases. Additionally, generating tokens simultaneously allows MJGD to return the optimal tokens for step . In contrast, vanilla greedy decoding selects each token without considering any future tokens. So MJGD is less prone to choosing local optima than greedy decoding.
Empirical evidence supports our claim. Figure 1 shows the output perplexity and downstream ROUGE-L scores of MJGD with set to a constant , where . Notice that setting is equivalent to vanilla greedy decoding. We use beam decoding to approximate the sampling from the joint distribution . We can see that when increases from 1 to 2, there is a significant improvement in perplexity. The perplexity continues to decrease when increases further. It confirms our claim that increasing will increase the output perplexity. So is a better approximation to than . Here , and for , and is the total number of steps of MJGD.
Moreover, the ROUGE-L score also improves with , supporting our claim that better perplexity reflects enhanced performance in downstream tasks. This demonstrates that MJGD not only achieves lower perplexity but also yields higher quality text generation, making it a superior decoding strategy compared to vanilla greedy decoding333Additional discussions comparing MJGD and greedy decoding are provided in the Appendix F..
3.2 Multi-Token Joint Speculative Decoding
Unfortunately, the computation cost of MJGD is infeasible in practice, since the time and space complexity to compute the joint distribution of tokens is . Inspired by speculative decoding and the facts that “even when a small model is an order of magnitude smaller than a large model, only a small fraction of the small model’s prediction deviate from those of the large model” [20, 17], we propose multi-token joint speculative decoding (MJSD), which exploits a small model to accelerate MJGD approximately. The core idea is to (1) use the joint distribution output by to approximate 444It is also valid to approximate with . Without loss of generality, we consider non-warped distribution in the illustration of MJSD. and generate draft tokens from , then (2) use the large model to validate draft tokens in parallel and accept the longest draft sub-sequence that passes verification, and (3) sample an additional token from the distribution of the large model without extra overhead to ensure at least one token is generated at each iteration. However, it is still infeasible to directly generate draft tokens from the joint distribution . So we propose to further approximate this process with beam decoding, which is an effective and efficient algorithm to generate sequences with high likelihood. In this way, MJSD reduces the number of runs of the large model to generate tokens, thus accelerating the inference in the same way as vanilla speculative decoding does. Algorithm 1 in the Appendix illustrates the pseudocode of MJSD algorithm.
![Refer to caption](x1.png)
Draft Tokens Verification
Figure 2 illustrates the verification process of MJSD. Let be the draft tokens generated by beam decoding with the small model. Since beam decoding is a widely recognized algorithm to generate sequences with high overall likelihood [18], it is reasonable to assume is large. Also, since beam decoding works in an autoregressive way, we can also assume that , is large. To approximate MJGD, for each step , MJSD needs to ensure the accepted tokens () also have high joint likelihood with the large model . So MJSD first computes the joint likelihood for . Then for each prefix sub-sequence , it passes verification if and only if , where is a pre-defined threshold. Notice that if , we have , which means . Therefore, our acceptance policy guarantees that when , the relative error is bounded. And if , it means the sub-sequence has higher likelihood in the large model, then it is reasonable to accept it. After verifying all the sub-sequences, MJSD accepts the longest draft sub-sequence that passes verification.
The verification step of MJSD ensures that the accepted tokens have a high joint likelihood with the large model. We have shown that selecting multiple tokens based on their joint likelihood lead to better output perplexity. Thus, MJSD is more effective than greedy decoding and vanilla speculative decoding. Furthermore, since MJSD accepts the longest draft sub-sequence with high likelihood, it can tolerate low-quality tokens as long as the joint likelihood is high. So at each iteration, MJSD can accept more draft tokens than vanilla speculative decoding, which results in better efficiency.
4 Energy Efficiency Analysis
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]:
(2) |
where denote the power (energy/second) of FLOPs and memory operations, and denote the time spent on these operations. When input batch size increases, increases until it reaches the power of maximum FLOPs, denoted as . Meanwhile, is irrelevant to the input batch size, as it only depends on the memory hardware.
To determine the relative magnitude relationship between and , we first point out the fact that GPU memory operations in LLM inference are dominated by accessing off-chip global memory, which consumes about 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 - more power than running maximum FLOPs [1]. So we have . Notice that 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 .
In addition, previous studies have shown that during LLM inference [20]. Therefore, the energy induced by memory operations, i.e., dominates . Since speculative decoding reduces 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 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 dominates . In the next section, we will show more direct evidence for the energy efficiency of MJSD and other speculative decoding algorithms.
Batch Size | Energy (J) | Energy/run (J) | Energy/Input (J) | Time (s) | Time/run (s) | Time/input(s) |
---|---|---|---|---|---|---|
1 | 42,450 | 14.1 | 14.1 | 1,129 | 0.376 | 0.376 |
2 | 49,621 | 16.5 | 8.26 | 1,191 | 0.397 | 0.198 |
4 | 53,325 | 17.7 | 4.43 | 1,178 | 0.392 | 0.098 |
8 | 59,210 | 19.7 | 2.46 | 1,211 | 0.403 | 0.050 |
16 | 74,058 | 24.7 | 1.54 | 1,255 | 0.418 | 0.026 |
5 Experiments
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 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 .
\hlineB3 | greedy | speculative | BiLD | Spectr | MJSD | ||
\hlineB3 CP | Llama-2 | speed (token/s) | 22.60.03 | 36.80.53 | 34.40.87 | 45.11.32 | 63.00.20 |
speed up | 1.000.00 | 1.630.02 | 1.520.04 | 1.990.06 | 2.780.01 | ||
energy (J/token) | 11.71.56 | 6.620.91 | 7.450.90 | 5.170.88 | 3.380.02 | ||
perplexity | 3.740.14 | 3.640.11 | 3.150.06 | 3.640.08 | 2.060.06 | ||
OPT | speed (token/s) | 22.40.48 | 33.82.47 | 31.51.87 | 38.02.20 | 55.80.30 | |
speed up | 1.000.00 | 1.510.08 | 1.410.06 | 1.700.06 | 2.500.05 | ||
energy (J/token) | 13.20.28 | 7.480.07 | 8.750.13 | 6.080.11 | 3.610.03 | ||
perplexity | 5.490.15 | 5.470.11 | 4.510.09 | 5.270.09 | 3.000.09 | ||
CA | Llama-2 | speed (token/s) | 22.00.15 | 31.60.35 | 28.80.20 | 27.70.29 | 44.10.25 |
speed up | 1.000.00 | 1.430.02 | 1.310.01 | 1.260.00 | 2.000.02 | ||
energy (J/token) | 11.20.16 | 6.980.15 | 7.990.15 | 7.200.08 | 4.720.03 | ||
perplexity | 2.110.01 | 2.130.03 | 1.950.03 | 2.150.01 | 1.880.05 | ||
OPT | speed (token/s) | 23.80.10 | 35.60.45 | 38.50.93 | 28.40.34 | 49.60.42 | |
speed up | 1.000.00 | 1.490.01 | 1.620.03 | 1.190.01 | 2.080.03 | ||
energy (J/token) | 11.30.22 | 5.740.11 | 5.120.06 | 6.240.11 | 4.030.02 | ||
perplexity | 3.280.06 | 3.320.10 | 2.600.06 | 3.160.06 | 2.070.03 | ||
CD | Llama-2 | speed (token/s) | 21.70.08 | 30.70.18 | 30.50.21 | 25.00.31 | 44.20.99 |
speed up | 1.000.00 | 1.410.00 | 1.410.01 | 1.150.01 | 2.040.05 | ||
energy (J/token) | 11.30.21 | 7.070.19 | 7.410.16 | 8.220.19 | 4.800.12 | ||
perplexity | 2.880.04 | 2.870.08 | 2.930.03 | 3.060.11 | 2.630.10 | ||
OPT | speed (token/s) | 23.30.81 | 31.70.91 | 30.90.80 | 23.70.40 | 43.60.33 | |
speed up | 1.000.00 | 1.360.04 | 1.320.04 | 1.010.02 | 1.870.02 | ||
energy (J/token) | 11.50.14 | 6.370.11 | 6.710.17 | 7.310.17 | 4.860.03 | ||
perplexity | 3.930.14 | 3.970.06 | 3.740.09 | 4.040.07 | 3.170.06 | ||
\hlineB3 |
5.1 Comparison with Baselines
Table 2 shows the main results of our experiments. The standard deviations in the table are computed by repeating the experiment four times. First, we compare the output perplexity of different algorithms. We can see that the perplexity of greedy decoding, vanilla speculative decoding, and Spectr decoding are close because they have identical sampling distributions for each token. Meanwhile, BiLD approximates the sampling distribution of greedy decoding but yields better perplexity. It is because we set a strict acceptance threshold for BiLD decoding, which lowers the acceptance rate but ensures every token has a high probability in the large model. It can be viewed as having an additional warping operation towards the sampling distribution. Nevertheless, there is a significant gap between the perplexity of MJSD and other baselines. On average, the output perplexity of MJSD is lower than that of BiLD. It provides strong evidence that MJSD has better effectiveness than existing speculative decoding methods that aim to have identical sampling distribution as greedy decoding.
Next, we compare the speeds of different algorithms. While the speeds of all speculative decoding baselines are close, MJSD is significantly faster. On average, MJSD is faster than greedy decoding, faster than vanilla speculative decoding, faster than BiLD, and faster than Spectr. The speed improvement is because MJSD has a larger average acceptance length at each iteration, as MJSD accepts the longest sub-sequence that passes verification. Figure 3 shows the average acceptance length of vanilla speculative decoding and MJSD. We can see that the acceptance length of MJSD is on average larger than that of vanilla speculative decoding.
Finally, we focus on energy efficiency. We can see that all speculative decoding algorithms have significantly better energy efficiency over greedy decoding. It supports our analysis that speculative decoding can improve the LLM inference energy efficiency by reducing the number of memory operations. In addition, we can see that the energy efficiency of MJSD is significantly better than that of all baseline algorithms. We believe it is also because of the longer acceptance sequence at each iteration. The number of memory operations is proportional to the number of LLM runs, which is inversely proportional to the average acceptance length. Since by always accepting the longest prefix that passes verification, the energy consumption of MJSD is on average smaller than that of vanilla speculative decoding, and smaller than that of greedy decoding.
![Refer to caption](extracted/5728708/figures/accepted_len.png)
5.2 Ablation Study
# beams | 2 | 4 | 6 | 8 | |
---|---|---|---|---|---|
speed (token/s) | 55.9 | 59.9 | 60.2 | 61.3 | |
Llama-2 | energy (J/token) | 2.43 | 2.25 | 2.22 | 2.20 |
perplexity | 2.44 | 2.12 | 2.14 | 2.10 | |
speed (token/s) | 51.0 | 54.1 | 54.3 | 55.9 | |
OPT | energy (J/token) | 2.50 | 2.32 | 2.36 | 2.30 |
perplexity | 3.63 | 3.16 | 3.42 | 3.19 |
5.2.1 Effects of Number of Beams
First, we investigate how the number of beams used in the beam decoding of the small model affects the inference performance. Table 3 shows the results. Increasing the number of beams improves the quality of the draft tokens, which not only improves the output perplexity but also increases the average acceptance length and hence leads to better efficiency. But we can see that the increment slows down when the number of beams is large enough. In addition, when the number of beams is too large, the inference cost of the small model will become too high.
5.2.2 Effects of Acceptance Threshold
![Refer to caption](extracted/5728708/figures/ppl_tau.png)
![Refer to caption](extracted/5728708/figures/speed_tau.png)
Next, we evaluate the effect of acceptance threshold . Intuitively, when we increase from 0 to 1, the acceptance criterion becomes more strict, the efficiency drops while the output perplexity increases. Surprisingly, this expectation is only partially correct. As shown in Figure 4, the efficiency indeed drops when increases. However, the perplexity increases when is close to 1. When , all the draft tokens are rejected, which makes MJSD equivalent to greedy decoding. Similarly, when is close to 1, the advantage of multi-token joint greedy decoding on effectiveness disappears, hence the perplexity becomes close to the perplexity of greedy decoding. Another surprising observation is that the perplexity of MJSD is good when . When , MJSD is equivalent to generating tokens using beam decoding with the small model, then generating an additional token with the large model. The fact that MJSD achieves good perplexity when can be explained by the fact that “even when a small model is an order of magnitude smaller than a large model, only a small fraction of the small model’s predictions deviate from those of the large model” [17, 20]. Moreover, when ranges from to , the performance of MJSD is relatively stable, suggesting that MJSD is not sensitive to the acceptance threshold.
6 Related Work
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 (). 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 . The large model rejects a token if its negative log-likelihood is larger than threshold . 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 . 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.
7 Conclusion
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.
8 Limitations and Impact Statements
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.
References
- [1] Tyler Allen and Rong Ge. Characterizing power and performance of gpu memory access. In 2016 4th International Workshop on Energy Efficient Supercomputing (E2SC), pages 46–53. IEEE, 2016.
- [2] Feifei Bear. Llmspeculativesampling. https://github.com/feifeibear/LLMSpeculativeSampling, 2024. Accessed: 2024-05-19.
- [3] Ning Bian, Hongyu Lin, Yaojie Lu, Xianpei Han, Le Sun, and Ben He. Chatalpaca: A multi-turn dialogue corpus based on alpaca instructions. https://github.com/cascip/ChatAlpaca, 2023.
- [4] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
- [5] Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023.
- [6] Jianmin Chen, Bin Li, Ying Zhang, Lu Peng, and Jih-kwon Peir. Tree structured analysis on gpu power study. In 2011 IEEE 29th International Conference on Computer Design (ICCD), pages 57–64. IEEE, 2011.
- [7] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1–113, 2023.
- [8] Alex de Vries. The growing energy footprint of artificial intelligence. Joule, 7(10):2191–2194, 2023.
- [9] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323, 2022.
- [10] Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint arXiv:1902.09574, 2019.
- [11] Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict: Parallel decoding of conditional masked language models. arXiv preprint arXiv:1904.09324, 2019.
- [12] Jiatao Gu, James Bradbury, Caiming Xiong, Victor OK Li, and Richard Socher. Non-autoregressive neural machine translation. arXiv preprint arXiv:1711.02281, 2017.
- [13] Junliang Guo, Linli Xu, and Enhong Chen. Jointly masked sequence-to-sequence model for non-autoregressive neural machine translation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 376–385, 2020.
- [14] Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D Lee, and Di He. Rest: Retrieval-based speculative decoding. arXiv preprint arXiv:2311.08252, 2023.
- [15] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015.
- [16] Norman P Jouppi, Doe Hyun Yoon, Matthew Ashcraft, Mark Gottscho, Thomas B Jablin, George Kurian, James Laudon, Sheng Li, Peter Ma, Xiaoyu Ma, et al. Ten lessons from three generations shaped google’s tpuv4i: Industrial product. In 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA), pages 1–14. IEEE, 2021.
- [17] Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. Speculative decoding with big little decoder. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
- [18] Rémi Leblond, Jean-Baptiste Alayrac, Laurent Sifre, Miruna Pislar, Jean-Baptiste Lespiau, Ioannis Antonoglou, Karen Simonyan, and Oriol Vinyals. Machine translation decoding beyond beam search. arXiv preprint arXiv:2104.05336, 2021.
- [19] Jason Lee, Elman Mansimov, and Kyunghyun Cho. Deterministic non-autoregressive neural sequence modeling by iterative refinement. arXiv preprint arXiv:1802.06901, 2018.
- [20] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pages 19274–19286. PMLR, 2023.
- [21] Zhuohan Li, Zi Lin, Di He, Fei Tian, Tao Qin, Liwei Wang, and Tie-Yan Liu. Hint-based training for non-autoregressive machine translation. arXiv preprint arXiv:1909.06708, 2019.
- [22] Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023.
- [23] Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, Ion Stoica, Zhijie Deng, Alvin Cheung, and Hao Zhang. Online speculative decoding. arXiv preprint arXiv:2310.07177, 2023.
- [24] Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rae Ying Yee Wong, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative llm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 1(2):4, 2023.
- [25] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485–5551, 2020.
- [26] Mohamed Rashad. Chatgpt-prompts, 2023.
- [27] Siddharth Samsi, Dan Zhao, Joseph McDonald, Baolin Li, Adam Michaleas, Michael Jones, William Bergeron, Jeremy Kepner, Devesh Tiwari, and Vijay Gadepally. From words to watts: Benchmarking the energy costs of large language model inference. In 2023 IEEE High Performance Extreme Computing Conference (HPEC), pages 1–9. IEEE, 2023.
- [28] Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. Advances in Neural Information Processing Systems, 33:20378–20389, 2020.
- [29] Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 35:17456–17472, 2022.
- [30] Abigail See, Peter J. Liu, and Christopher D. Manning. Get to the point: Summarization with pointer-generator networks. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1073–1083, Vancouver, Canada, July 2017. Association for Computational Linguistics.
- [31] Zhiqing Sun, Zhuohan Li, Haoqing Wang, Di He, Zi Lin, and Zhihong Deng. Fast structured decoding for sequence models. Advances in Neural Information Processing Systems, 32, 2019.
- [32] Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, and Felix Yu. Spectr: Fast speculative decoding via optimal transport. arXiv preprint arXiv:2310.15141, 2023.
- [33] Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239, 2022.
- [34] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
- [35] Yiren Wang, Fei Tian, Di He, Tao Qin, ChengXiang Zhai, and Tie-Yan Liu. Non-autoregressive machine translation with auxiliary regularization. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 5377–5384, 2019.
- [36] Yisheng Xiao, Lijun Wu, Junliang Guo, Juntao Li, Min Zhang, Tao Qin, and Tie-yan Liu. A survey on non-autoregressive generation for neural machine translation and beyond. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
- [37] Zeyu Yang, Karel Adamek, and Wesley Armour. Part-time power measurements: nvidia-smi’s lack of attention. arXiv preprint arXiv:2312.02741, 2023.
- [38] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
Appendix A Pseudocode of MJSD
See Algorithm 1.
Dataset | Task | # Test Input | Avg. Input Len |
---|---|---|---|
ChatGPT-Prompt | Instruction | 360 | 25.2 |
ChatAlpaca | Chat | 43.081 | 277.7 |
CNNDM | Summarization | 11,490 | 3,967.1 |
Appendix B Energy Consumption Measurement
We use the command "nvidia-smi --query-gpu=power.draw --format=csv" to get GPU power every second, and sum them up as the total energy consumption. We use average energy consumption per token to measure energy efficiency. There is a recent study pointing out the measurement error using nvidia-smi [37]. We follow the three principles proposed in [37] to minimize the error.
Appendix C Configuration
The experiments are conducted on a machine with 1 Nvidia L40 GPU (48 GB), 4 CPUs, and 50 GB main memory, using a batch size of 1, which is common for online serving [29]. We set the maximum running time to be an hour for each baseline. We use average tokens/second to measure the inference speed and use perplexity (exponentiated average negative log-likelihood) based on the probability of the large model to measure the output quality. Because different methods might finish different numbers of inputs, we only calculate the perplexity of the first inputs, where is the number of inputs finished by greedy decoding. We use average energy consumption per token to measure energy efficiency. The details of energy measurement are illustrated in the Appendix.
Appendix D Hyper-parameter Details
In the experiments, we follow the default settings in [2] to warp the sampling distribution and with the following steps, which are the default warpping operations in a public speculative decoding implementation.
-
1.
Keep the probabilities of top 20 tokens unchanged, and set the probabilities of other tokens to 0, then normalize the distribution.
-
2.
Sort the tokens based on their distributions descendingly. Keep the first tokens such that their cumulative probabilities is larger than 0.9, while set the probabilities of other tokens to 0.
For different methods, we choose their hyper-parameters by using a small validation set. We select the set of hyper-parameters that make the corresponding method have best output perplexity. Table 5 shows the hyper-parameters used in the experiments.
Llama,CP | OPT,CP | Llama,CA | OPT,CA | Llama,D | OPT,CD | ||
speculative | step len | 4 | 4 | 4 | 4 | 4 | 4 |
Spectr | step len | 4 | 4 | 4 | 4 | 4 | 4 |
num seq | 4 | 4 | 2 | 2 | 4 | 2 | |
BiLD | step len | 10 | 10 | 10 | 10 | 10 | 10 |
fallback thres | 0.9 | 0.9 | 0.9 | 0.3 | 0.9 | 0.3 | |
rollback thres | 2 | 2 | 1 | 2 | 3 | 2 | |
MJSD | step len | 4 | 4 | 4 | 4 | 4 | 4 |
num beams | 8 | 8 | 8 | 8 | 8 | 8 | |
acc/rej thres | 0.1 | 0.1 | 0.1 | 0.1 | 0.1 | 0.1 |
Appendix E License of Datasets and Models
Datasets:
-
•
ChatGPT-Prompts: Non (https://huggingface.co/datasets/MohamedRashad/ChatGPT-prompts)
-
•
ChatAlpaca: Apache-2.0 License
-
•
CNN Dailymail: Apache-2.0 License
Models
-
•
OPT-125M and OPT-13B: Model License (https://github.com/facebookresearch/metaseq/blob/main/projects/OPT/MODEL_LICENSE.md)
-
•
Llama-68M: Apache-2.0 License
-
•
Llama-2-13B: Llama-2 Community License Agreement
Codes
-
•
LLMSpeculativeSampling (https://github.com/feifeibear/LLMSpeculativeSampling)
Appendix F More Discussion on Multi-Token Joint Greedy Decoding
It is not trivial to prove that MJGD is better than vanilla greedy decoding. We might tend to show that given the same and , the tokens generated by MJGD has higher likelihood than the tokens generated by calling greedy decoding times. However, at step (), the previously generated tokens are different for the two decoding algorithms, making it difficult to compare the overall perplexity.
To explain why MJGD is more effective, let us consider an optimal algorithm that generates next tokens at one step based on a score function:
(3) |
which is the log-likelihood of the optimal future sequence with prefix and the generated tokens. The optimal greedy decoding chooses that maximizes the score function and is guaranteed to return the sequence with the optimal perplexity.
Both MJGD and vanilla greedy decoding can be viewed as approximations to the optimal greedy decoding by replacing with oder that estimates the values of . Specifically, since MJGD selects the tokens with the largest joint likelihood, we have where is an arbitrary constant. Notice that adding a constant to changes the error of the estimation but does not change the behavior of the algorithm.
On the other hand, greedy decoding generates the tokens one by one. At the last token, greedy decoding selects the token with the largest , which is equivalent to selecting the tokens with the largest provided that are the same tokens generated by previous steps. So for greedy decoding, ( is also an arbitrary constant) if and only if for . Otherwise, the score of has to be smaller than any possible value of . Notice that the lower bound of is , so if and only if such that .
Let . By setting , for satisfying that , , we have
Otherwise, we have
where . Therefore,
(4) |
The derivation above provides a heuristic explanation to show that MJGD is a closer approximation to the optimal greedy decoding than vanilla greedy decoding. Hence, its output perplexity is generally better.