Optimized Multi-Token Joint Decoding With Auxiliary Model for LLM Inference

Authors: Zongyue Qin, Ziniu Hu, Zifan He, Neha Prakriya, Jason Cong, Yizhou Sun
Primary area: foundation or frontier models, including LLMs
Keywords: LLM Inference, Speculative Decoding

Abstract:
Large language models (LLMs) have achieved remarkable success across diverse tasks, yet their inference processes are hindered by substantial time and energy demands due to single-token generation at each decoding step. While previous methods such as speculative decoding mitigate these inefficiencies by producing multiple tokens per step, each token is still generated by its single-token distribution,
thereby enhancing speed without improving effectiveness. In contrast, our work simultaneously enhances inference speed and improves the output effectiveness. We consider multi-token joint decoding (MTJD), which generates multiple tokens from their joint distribution at each iteration, theoretically reducing perplexity and enhancing task performance. However, MTJD suffers from the high cost of sampling from the joint distribution of multiple tokens. Inspired by speculative decoding, we introduce multi-token assisted decoding (MTAD), a novel framework designed to accelerate MTJD. MTAD leverages a smaller auxiliary model to approximate the joint distribution of a larger model, incorporating a verification mechanism that not only ensures the accuracy of this approximation, but also improves the
decoding efficiency over conventional speculative decoding. Theoretically, we demonstrate that MTAD closely approximates exact MTJD with bounded error. Empirical evaluations using Llama-2 and OPT models ranging from 13B to 70B parameters across various tasks reveal that MTAD reduces perplexity by 21.2% and improves downstream performance compared to standard single-token sampling.
Furthermore, MTAD achieves a 1.42× speed-up and consumes 1.54× less energy than conventional speculative decoding methods. These results highlight MTAD’s ability to make multi-token joint decoding both effective and efficient, promoting more sustainable and high-performance deployment of LLMs.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730189448928

thread_type:
official_review

<thread_object>
creator: reviewer_ds8n

<content>
summary:
Considering the inference efficiency, this paper proposes multi-token joint decoding (MTJD), which generates multiple tokens from their joint distribution at each iteration. More importantly, this paper theoretically proves that their method can reduce perplexity and improve task performance compared to single token decoding. Besides, to mitigate the high cost of sampling from the joint distribution of multiple tokens, it introduce a small model to assist the decoding of a larger model. Empirical evaluations reveal that the method reduces perplexity
by 21.2% and improves downstream performance compared to standard single-token sampling. It also achieves a 1.42× speed-up and consumes 1.54× less energy than conventional speculative decoding methods.

strengths:
1. This paper theoretically proves that multi-token joint decoding (MTJD) can reduce perplexity and improve task performance compared to single token decoding.
2. This paper introduces multi-token assisted decoding (MTAD), a novel framework designed to accelerate multi-token joint decoding. Theoretically, this paper demonstrate that MTAD closely approximates exact MTJD with bounded error.
3. These results highlight MTAD’s ability to make multi-token joint decoding both effective and efficient.

weaknesses:
1. This paper primarily conducted experiments on the OPT series and Llama 2 series models. However, the performance of these large models has already fallen far behind the recent Llama 3.1 and Llama 3.2. Could you provide more results on the latest models to demonstrate the generality of the results?

2. The previous methods compared in this paper were all published last year. It would be beneficial to compare with some recently published methods, such as MEDUSA-2 and "Better & Faster Large Language Models via Multi-token Prediction", to prove the effectiveness of the results.

3. In this paper, a major point is that multi-token joint decoding can reduce perplexity and improve task performance. However, the experimental results do not reflect related analytical experiments. Adding relevant experiments could enhance the self-consistency of the method.

questions:
Please refer to the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730671118760

thread_type:
official_review

<thread_object>
creator: reviewer_v4l1

<content>
summary:
This paper focuses on improving the inference speed of large language models (LLMs). The authors propose a variant of speculative sampling that enhances both inference speed and output quality. Specifically, they explore multi-token joint decoding (MTJD), which generates multiple tokens from their joint distribution at each iteration, theoretically reducing perplexity and improving task performance. Building on this, they incorporate speculative sampling to approximate multi-token joint decoding and design an efficient algorithm, multi-token assistant decoding (MTAD), which achieves bounded error with respect to MTJD. Empirical evaluations demonstrate the effectiveness of the proposed method in terms of speed and perplexity.

strengths:
- The writing is good and easy to follow.
- The addressed problem is of great importance to the community.
- A new metric is provided to evaluate the effectiveness of speculative sampling, which is insightful.
- Strong experimental results:
  - Evaluated across a range of datasets
  - Shows significant improvements in terms of perplexity, speed, and energy efficiency.

weaknesses:
- In line 114, the authors claim, “As shown in Table 1, lower perplexity correlates with improved downstream performance, even in one of today’s largest models.” This assertion is questionable. The results in Table 1 are based exclusively on experiments with the Spider dataset, but in many cases, lower perplexity does not inherently lead to better downstream performance. For example, in tasks like machine translation and abstractive summarization, a lower perplexity score sometimes fails to correlate with improved downstream performance. This also can be validated by the right plot in Figure 1.
- Lack of guarantee on quality: In vanilla speculative sampling and multi-draft speculative sampling, one of their advantages is the property of being lossless. This ensures that the quality of output is maintained, which further accelerates LLM inference. However, the proposed method does not have such guarantee. 
- The verification process in this work is similar to that in [1], as both aim to verify a sub-sequence of tokens in a single step rather than one token at a time. However, this study does not include a comparison between MTAD and the method in [1]. The author should further compare their work with it to further improve the quality.
- The experiments also have some limitations:
  - This paper places too much emphasis on perplexity. Since a lower perplexity score does not always correlate with improved downstream performance, the authors should provide a more comprehensive evaluation of downstream performance. I noticed that Table 6 contains some comparisons on downstream performance, but this is insufficient. I encourage the authors to add a "downstream performance" metric to Table 4 so that readers can better understand the effectiveness of the proposed method.
  - I still do not fully understand why the proposed method significantly outperformed multi-draft speculative sampling methods such as Specter and SpecInfer. I suspect this is because (1) the selected baseline is weak, and (2) the configurations are not fair. For example, in Table 13, the number of sequences in Specter and SpecInfer does not match the number of beams in MTAD. I encourage the authors to do the following: (1) select more recent multi-draft speculative sampling methods for comparison, and (2) ensure consistent configurations.


[1] Sun, et al. Optimal Block-Level Draft Verification for Accelerating Speculative Decoding.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730709895669

thread_type:
official_review

<thread_object>
creator: reviewer_h4x4

<content>
summary:
This paper introduces a new framework named multi-token assisted decoding (MTAD) by combining speculative decoding with multi-token joint decoding (MTJD). MTAD generates high quality draft tokens by estimating the joint distribution of the target model using draft model, while improving both the speed and quality of inference.

strengths:
1. The proposed MTAD guides the generation and verification of draft tokens based on joint probabilities, providing a new idea for speculative decoding.

2. This paper analyzes the speculative decoding technique from an energy perspective for the first time and verifies the energy efficiency of speculative decoding. 

3. In addition to the experimental analysis, the authors also theoretically demonstrate the bounded error as well as the effectiveness of the MTAD algorithm.

weaknesses:
1. Despite generating texts with lower perplexity, MTAD is still a lossy speculative decoding method. It is unfair to compare MTAD with lossless speculative decoding methods. It might be useful to add some comparisons with other lossy verification methods (e.g. Medusa's Typical Sampling) to verify the superiority of MTAD.

2. The main experiments in this paper use perplexity as an evaluation criterion for performance, and it is clear that MTAD based on PPL verification will have advantages. More experiments on downstream tasks are needed to verify whether MTAD can really improve the quality of text, since PPL and metrics are not strictly positively correlated.

questions:
1. The threshold for the verification phase of vanilla speculative decoding is sampled from U(0, 1) to ensure losslessness, and MTAD uses a fixed threshold (e.g. 0.1). I'm curious to see if the same speedup can be achieved with the vanilla speculative decoding method using a fixed threshold.

2. The impact from loose thresholds has little effect with PPL and LM-Judger metrics, some more rigorous evaluation (e.g., Humaneval's Pass@K) is needed to validate the methodology.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730739299347

thread_type:
official_review

<thread_object>
creator: reviewer_rq6h

<content>
summary:
The paper proposed multi-token assisted decoding by combining multi-token joint decoding and speculative decoding (SpD) to improve the quality of the tokens and inference speed.

strengths:
The paper solves an interesting problem of sampling from a joint distribution which is challenging and proposes an alternate approximate method based on draft model from Speculative Decoding (SpD) 

It uses ideas from SpD to also boost the inference speed.

Paper uses an auxiliary model to predict joint distribution of multiple tokens w/ beam sampling to sample the tokens fast. 

The power analysis provides some insight into the importance of using SpD

Authors provide a robust set of experiments

weaknesses:
the experiment results can be re-checked, my concern is that lower PPL not always results in better quality, it could also be that the draft model is generating repeated tokens which the target model is accepting, given that there's no guarantee that the draft tokens accepted will follow the target distribution. 

the reduction in energy consumption is a by-product of SpD and could have been mentioned in a sub-section rather a full section. 
The acceptance rate of draft model should affect both the token-rate and the energy consumption.

No mention of draft model or target model training and if they are different from the baseline algorithms draft and target model in the experiments section.

questions:
target model (larger model) is fine-tuned on chat-GPT data for proposed method, is the same fine-tuned model used for other Speculative Decoding (SpD) methods as well? 

What about the draft model? is that fine-tuned as well or used as is? 

Can qualitative generations be shown in the paper? There are cases when lower PPL is observed during SpD as the draft and target model are both generating and accepting repeated tokens.

Did the author also check performances (PPL, speed-up) when using a larger draft model? 

From table 4., it seems like for past SpD methods there is no correlation between high speed (token/s) will lead to low energy (Joule/s), but that doesn't seem to make sense, as both high speed and low energy should be related to high acceptance rate which leads to small number of target model calls. Any comment/thoughts on this? 

Can the authors also add plots on acceptance threshold vs acceptance rate? Given the high token-rates shown in Table 4., I am curious to see how the acceptance rates look like.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732162760022

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common reply about perplexity and downstream performance

comment:
We acknowledge the reviewers’ concerns regarding whether MTAD’s lower perplexity also translates to improved output quality. In our original submission (Table 6), we demonstrated MTAD’s superior output quality using multiple metrics across diverse datasets with the Llama-2-13B model. **To address these concerns further, we provide additional justification and evaluation below.**

**Multi-Token Joint Decoding (MTJD) has better output quality**  
While most decoding approaches focus on inference speed up, we want to design approaches that can also improve inference quality. We propose multi-token joint decoding (MTJD Section 3.1) to achieve the goal, due to its capability to achieve lower perplexity and higher likelihood than single-token multinomial sampling. To validate that MTJD indeed improves output quality, we evaluate MTJD (k=4) and multinomial sampling on Spider, MTBench, and HumanEval (Pass@1) using the Llama-3 series models. We follow the same way introduced in Section 3.1 to implement MTJD. The results, where higher scores indicate better performance, show a clear advantage for MTJD in terms of output quality. 

|  |  | Llama-3-8B and Llama-3-1B |  |  | Llama-3-8B-Instruct and Llama-3-1B-Instruct |  |  |
| :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- |
| dataset |  | Spider | MTBench | HumanEval | Spider | MTBench | HumanEval |
| Single-token multinomial sampling | score | 22.0 | 3.40 | 15.9 | 36.0 | 4.11 | 28.0 |
|  | ppl | 2.58 | 2.40 | 2.09 | 2.23 | 1.91 | 1.85 |
| Multi-token joint sampling | score | 52.5 | 3.77 | 36.6 | 60.5 | 4.40 | 49.4 |
|  | ppl | 1.16 | 1.32 | 1.26 | 1.18 | 1.27 | 1.15 |

Since MTJD is computationally expensive (slower than greedy decoding), we further propose MTAD to approximate MTJD, inheriting its benefits while achieving significantly better efficiency. The performance of MTAD is discussed below.

**Comprehensive Evaluation of MTAD**  
To further validate MTAD’s effectiveness, we conducted additional evaluations, addressing reviewers’ requests: (1) we use Llama-3 models (Llama-3.1-8B and Llama-3.1-8B-Instruct as target models, and Llama-3.2-1B and Llama-3.2-1B-Instruct as draft models); (2) we include three datasets: Spider, HumanEval, and MT-Bench, to include rigorous metrics and a variety of tasks; (3) we add Medusa’s \[1\] decoding algorithm, typical sampling, as an additional baseline. **The results below show that MTAD consistently achieves the best output quality and the fastest speed across all datasets and models.**

|  |  |  | Lossy Speculative Decoding |  | Lossless Speculative Decoding |  |  | | Ours |
| :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- |:----|
|  |  |  | BiLD | Typical (Medusa) | spec | spectr | specinfer|MCSS | MTAD |
| Spider | Llama-3-Instruct | tokens/s | 20.1 | 22.3 | 19.6 | 22.4 | 21.1 | 21.7| **23.5** |
|  |  | J/token | 10.2 | 9.5 | 10.5 | 9.6 | 10.2 |10.0 |**9.2**|
|  |  | Acc | 35.0 | 42.0 | 36.0 | 35.5 | 37.0 |35.0|**44.0**|
|  | Llama-3 | tokens/s | 23.3 | 32.3 | 31.1 | 32.1 | 32.6 |32.7| **33.3**|
|  |  | J/token | 8.2 | 7.9 | 7.5 | **7.1** | 8.1 | 8.0|7.8 |
|  |  | Acc | 30.5 | 29.5 | 21.5 | 23.0 | 21.5 |24.0| **35.0** |
| MT-Behcn | Llama-3-Instruct | tokens/s | 25.9 | 23.4 | 26.0 | 26.2 | 26.3 |26.8| **29.8** |
|  |  | J/token | 10.8 | 12.2 | 10.0 | 9.9 | 10.0 |9.9| **9.2** |
|  |  | score | 4.15 | 4.26 | 4.1 | 4.11 | 4.01 |4.02| **4.40** |
|  | Llama-3 | tokens/s | 24.5 | 22.3 | 24.1 | 24.5 | 24.5 |25.7| **28.2** |
|  |  | J/token | 11.5 | 12.4 | 11.0 | 11.6 | 11.7 | 11.1|**10.0** |
|  |  | score | 3.41 | 3.24 | 3.39 | 3.41 | 3.35 |3.36| **3.75** |
| HumanEval | Llama-3-Instruct | tokens/s | 17.4 | 21.7 | 22.2 | 23.8 | 22.8 |23.7| **24.8** |
|  |  | J/token | 10.0 | 8.1 | 7.8 | 7.8 | 7.9 |7.8| **7.6** |
|  |  | pass@1 | 37.8 | 35.9 | 32.9 | 32.9 | 31.0 |32.0|**38.4**|
|  | Llama-3 | tokens/s | 19.6 | 22.5 | 22.2 | 24.4 | 22.5 |23.8|**25.6**|
|  |  | J/token | 9.7 | 8.9 | 8.9 | 8.9 | 8.1 |7.9|**7.6**|
|  |  | pass@1 | 19.5 | 20.0 | 15.9 | 16.0 | 17.7 |17.0| **22.0** |

**Visualization of relationship between perplexity and downstream performance**

To further illustrate the relationship between perplexity and downstream performance, we present a scatter plot (Figure 4 in the Appendix C.4). The plot shows the correlation between **relative downstream scores** (normalized by the score of multinomial sampling) and **relative perplexity** (normalized by the perplexity of multinomial sampling) across 7 decoding algorithms, 3 datasets, and 2 model configurations. The results confirm that lower perplexity generally correlates with higher output quality.

We sincerely hope these additional results and analyses could address the reviewers’ concerns and demonstrate the robustness, efficiency, and quality improvements achieved by MTAD.

\[1\] Cai, Tianle, et al. "Medusa: Simple llm inference acceleration framework with multiple decoding heads." *arXiv preprint arXiv:2401.10774* (2024).
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces Multi-Token Assisted Decoding (MTAD), a speculative decoding framework combining multi-token joint decoding (MTJD) with speculative decoding (SpD).  The proposed framework demonstrates significant gains in perplexity (21.2% reduction), task performance, speed-up (1.42×), and energy efficiency (1.54×) across multiple datasets.

**Strengths** (1) The paper improves speculative decoding with multi-token joint decoding, which advances beyond single-token-based methods. A theoretical analysis of bounded error provides additional insight into the method. (2) Comprehensive evaluation of multiple datasets and tasks.

**Weaknesses** (1) Limited comparison with recent baselines and qualitative analysis (2) Some reviewers raised concerns about the relationship between perplexity and quality

**Decision** Considering the novelty of the proposed methods, and how all the reviewers found the rebuttal addressing their original concerns, I am leaning toward acceptance.