Mixture of Attentions For Speculative Decoding

Authors: Matthieu Zimmer, Milan Gritta, Gerasimos Lampouras, Haitham Bou Ammar, Jun Wang
Primary area: generative models
Keywords: large language models, speculative decoding, EAGLE

Abstract:
The growth in the number of parameters of Large Language Models (LLMs) has led to a significant surge in computational requirements, making them challenging and costly to deploy.
Speculative decoding (SD) leverages smaller models to efficiently propose future tokens, which are then verified by the LLM in parallel.
Small models that utilise activations from the LLM currently achieve the fastest decoding speeds.
However, we identify several limitations of SD models including the lack of on-policyness during training and partial observability. 
To address these shortcomings, we propose a more grounded architecture for small models by introducing a Mixture of Attentions for SD.
Our novel architecture can be applied in two scenarios: a conventional single device deployment and a novel client-server deployment where the small model is hosted on a consumer device and the LLM on a server.
In a single-device scenario, we demonstrate state-of-the-art speedups improving EAGLE-2 by 9.5% and its acceptance length by 25%.
In a client-server setting, our experiments demonstrate: 1) state-of-the-art latencies with minimal calls to the server for different network conditions, and 2) in the event of a complete disconnection, our approach can maintain higher accuracy compared to other SD methods and demonstrates advantages over API calls to LLMs, which would otherwise be unable to continue the generation process.

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

Reviews and discussion:
<thread>
creation date:
1729909189348

thread_type:
official_review

<thread_object>
creator: reviewer_sr5b

<content>
summary:
The paper proposes an improved architecture for speculative decoding, which incorporates Layer Self-Attention and Cross-Attention mechanisms to address common limitations like partial observability and lack of on-policy training. Experimental results demonstrate that this approach achieves state-of-the-art speed and accuracy in single-device and client-server deployments, maintaining high accuracy even during disconnections. The limitation analysis of SD in the paper provides insights into this field.

strengths:
(1) The limitation analysis of SD in the paper is very interesting, it links the limitation of SD (e.g., the SD may still give different predictions compared to only using the large model even when the previous tokens are all accepted ) with theoretical analysis, i.e., partial observability. This may provide further insights of optimizing SD algorithms.

(2) It is a very good point that the paper design and evaluate the scenario that the small model is on a resource-limited client and the large model is on the server. It is a very realistic scenario to use SD.

(3) The paper uses SOTA models such as Llama-3 for evaluation. Also the benchmarks used (e.g., HumanEval, MT-Bench) are challenging enough for evaluation.

weaknesses:
(1) The link between the limitation of SD and the proposed attention-based method is not clear enough. Why Layer Self-Attention can solve partial observability? Further intuitive explanations or theoretical analyses are needed.

(2) Authors mention that their work is based on a SOTA SD method named EAGLE. Then my concern is, is the proposed algorithm transferable to other SD algorithms? For example, can the algorithm be applied to non-self-draft SD? The point here is that, if the algorithm can be applied to most of the SD frameworks, there will be more contribution. Otherwise, it's just an optimization of one of the previous works.

(3) No accuracy is displayed in the experiment section. I know that this is common in SD papers. However, when I was trying SD codes, I found that the SD performance is usually not as good as the large model, which is the purpose of SD. Thus, because this paper aims to solve problems such as partial observability, will it also increase the performance of SD, in addition to the efficiency? I believe experiments related to this point will make the paper more convincing.

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

</thread>
<thread>
creation date:
1730367673908

thread_type:
official_review

<thread_object>
creator: reviewer_wo7q

<content>
summary:
This paper focuses on an LLM acceleration technique called Speculative Decoding, which leverages efficient models (smaller but less capable) to draft future tokens, which are verified by the LLM (more capable but much less efficient) in parallel. In particular, it addresses the limitations of previous methods by proposing a Mixture of Attentions architecture on top of a prior work EAGLE to improve its performance. They demonstrate the effectiveness of their approach in both a single-device setting and a novel client-server setting, achieving speedups and improved accuracy. They also present a framework for using LLMs on edge devices, allowing for offline text generation with minimal dependence on a server.

strengths:
- The work improves above EAGLE-2 and seems to achieve state-of-the-art results.
- The work provides a good background on speculative decoding
- The work proposes an interesting client-server setup that fits well with the speculative decoding technique

weaknesses:
- The work lacks an overall view and clear statements that can improve readability.
    - Method intuition: the method section only lays out the information of each component but does not provide an overall view of the proposed method as well as motivating intuitions for each design. The necessary intuitive descriptions are also not found in the appendix.
    - Experiment result: the work only compares to one prior work, EAGLE-2, as a baseline, but did not provide information on how well EAGLE-2 compares with other prior works.
    - the hyperparameter N is used many times. Giving it a name can help readability.
- This work overlooks overall performance metrics, concentrating instead on metrics specific to the speculative decoding framework, such as acceptance length. This focus may inadvertently encourage adversarial scenarios where the smaller model aims merely to deceive the larger model into accepting its outputs rather than genuinely enhancing result quality.
- The work neglects analysis of the time/computational complexity of its method. 
- The ablation study compares different hyperparameters but does not ablate the other components. Thus, the importance and contributions of the designs in Section 3.1 and Section 3.3 is not provided.

questions:
(also see weaknesses)
- The experiment on N is arbitrarily set to 0, 1, 3. Why not continuously evaluate 0,1,2.... to the number of layers in the large model?
- In the client-server mode with N > 0, are the last layers of the large model copied onto the client side?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730719913797

thread_type:
official_review

<thread_object>
creator: reviewer_2fbk

<content>
summary:
This paper proposes a method, Mixture of Attentions for SD, which improves upon the standard SD approach that is used to increase LLM efficiency. The improvements are specifically targeting the partial observability and lack of on-policyness problems of traditional SD. The Mixture of Attentions method is shown to lead to improvements for a single device, but the paper also shows its efficiency/accuracy benefits to a client-server setting.

strengths:
The organization and flow of the paper is very good. The background section is particularly thorough and helpful.

The problem is well-explained (e.g. partial observability and lack of on-policyness are both detailed when explaining the methodology) so it is made clear what exactly the Mixture of Attentions method is aiming to solve. Additionally, the related work is well-addressed. It is clear exactly how this work is different from prior solutions.

It is great that the client-server scenario is tested in a practical setting with different devices (having different resource capacities) and the devices have distance between them. This setting is not only realistic, but it is also well-explained in the text.

There is good theoretical support in the Methodology section, which provides additional justification for the proposed method being superior to the current SOTA (EAGLE, Medusa).

The thorough experiment detail (particularly in the appendix) makes the method highly reproducible.

weaknesses:
The experimentation is very narrow, especially since it only focuses on one model architecture and the improvements over EAGLE seem relatively small and inconsistent. It is therefore not convincing that this method would be effective more generally.

It is not very clear why this problem/contribution is important. The paper would be stronger if the method was motivated by some real-world example where SD may be used, but would lead to significant problems that Mixture of Attentions would mitigate. I understand that the computational requirements of LLMs is an issue, but the introduction could do a better job of explaining why SD should be focused on as a solution for the computational expense and therefore is important to build improvements for. It is also difficult to understand how this work can be have a broader impact or inspire future work. The future work that is suggested at the end of the conclusion seems very specific and narrow. Essentially, the contribution just seems very narrow.

questions:
Why did you choose to only experiment with LLama3-8B-Instruct? It is good that there is justification for only comparing to EAGLE (and not Medusa or Tandem Transformers), but there is no justification for your LLama model choice. How do you think your method would work with other architectures?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731052037962

thread_type:
official_review

<thread_object>
creator: reviewer_eflf

<content>
summary:
The method addresses two key limitations of existing SD approaches, (1) partial observability and (2) lack of on-policyness, by incorporating Layer Self-Attention (LSA) and Cross-Attention (CA).

strengths:
- The client-server framework with the ability to handle disconnections positions the approach as a practical advancement for deploying LLMs.

- The introduction of LSA and CA layers to mitigate partial observability and improve on-policyness makes sense.

weaknesses:
1. The paper does not thoroughly justify the choice of parameter configurations and its training in its experiments.  As discussed in the Yi et al. (2024), the training dataset and the choice of number of parameters can significantly affect the SD performance, but this paper does not [A].

[A] Yi et al., 2024. Towards Fast Multilingual LLM Inference: Speculative Decoding and Specialized Drafters, EMNLP 2024-main.

2. Discussions for the memory-bound nature of LLM is required in the paper.

3. The effectiveness of the proposed method on models having 3B~13B parameters is unclear. Current experiments focus on relatively smaller models, and the results may not hold for state-of-the-art LLMs, which typically exhibit different scaling dynamics and memory behavior.

4. Typo? line 466.

5. Parallel to Medusa, actually there are concurrent works regrading non-autoregressive heads for SD. It would be good to put discussions for those areas. 

- Gloeckle et al. (2024), Better & Faster Large Language Models via Multi-token Prediction.

- Stern et al. (2024), Blockwise Parallel Decoding for Deep Autoregressive Models

- Kim et al. (2024), Accelerating Blockwise Parallel Language Models with Draft Refinement. (https://openreview.net/forum?id=KT6F5Sw0eg)

questions:
See Weakness.

Will update the score after looking at the results of Weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732219942558

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We would like to answer a shared concern about the possible loss in generation quality with speculative decoding methods (reviewer Wo7q.Q4 and Sr5b.Q5). To do so, we ran several experiments to assess the quality of the generated responses using greedy decoding. We focused on 3 datasets from SpecBench (HumanEval, GSM8K and CNN/DM) that do not require access to proprietary models/APIs for evaluation (llm-as-a-judge).


| Vanilla decoding | HumanEval (pass@1) | GSM8K (accuracy) | CNN/DM (Rouge-L f-score) |
| ----------- | :----: | :----: | :----: |
| Llama-8B | 62.5% | 80% | 0.3071 |

| Speculative Decoding | HumanEval (pass@1) | GSM8K (accuracy) | CNN/DM (Rouge-L f-score) |
| ----------- | :----: | :----: | :----: |
| Ours (N=3) | 62.5% | 80% | 0.3053 |
| Ours (N=1)   | 62.5% | 81.25% | 0.3068 |
| Ours (N=0)   | 62.5%  | 80% | 0.3070 |
| EAGLE-2   | 62.5%  | 81.25% | 0.3062 |
| EAGLE-2 off   | 62.5%  | 80% | 0.3056 |
| Independent 1.7B   | 62.5%  | 80% | 0.3067 |
| Independent 1.3B   | 62.5%  | 80% | 0.3064 |

As you can observe in the tables above, the pass@1 on HumanEval is the same across all methods. The accuracy on GSM8K actually improves w.r.t the base model on one question for Ours (N=1) and EAGLE-2. Finally, the ROUGE scores are also extremely similar, leading us to conclude that any differences to the base model are negligible and almost certainly appear due to using float16 (see our answer to Sr5b.Q5 for more details).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732541009755

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Qwen2.5-3B

comment:
Dear reviewers,

Thank you very much for your suggestions to make the paper stronger. 
We would like to answer a shared concern about the limited range of models and sizes that we tested (reviewer eFLF.Q3 and 2FbK.Q3).
We trained 4 additional small models on the Ultrachat dataset to accelerate Qwen2.5 3B.
EAGLE recommends to use one decoder layer of the big LLM to define the size of the small LM, which leads to a trainable size of 80M parameters.
We kept the shared "embed_tokens/LM_head" layer frozen.

| Small Model | Total size | Trainable size   | MT-bench    |     | HumanEval |          | GSM8K    |          | Alpaca   |          | CNN/DM   |          | Natural Ques. |          | Mean     |          |
|--------------|-------|-----------|----------|----------|-----------|----------|----------|----------|----------|----------|----------|----------|---------------|----------|----------|----------|
| | | | Speedup  | $\tau$   | Speedup   | $\tau$   | Speedup  | $\tau$   | Speedup  | $\tau$   | Speedup  | $\tau$   | Speedup       | $\tau$   | Speedup  | $\tau$   |
| Ours (N=1)   | 0.48B  | 80M      | 1.7     | **3.61**     | 2.08    | **4.5**     | 1.5     | **3.31**     | 1.89 | **4.02**     | 1.78 | **3.81**     | 1.71      | **3.7**     | 1.76 | **3.82**     |
| Ours (N=0)   | 0.4B  | 80M      | **1.77**     | 3.56     | **2.16**    | 4.45     | **1.58**     | 3.24     | **1.89** | 3.82     | **1.83** | 3.71     | **1.69**      | 3.48     | **1.82** | 3.70     |
| EAGLE-2      | 0.4B  | 80M      | 1.65     | 3.13     | 1.88      | 3.57     | 1.5     | 3    | 1.78     | 3.35     | 1.7     | 3.14     | 1.63          | 3.05     | 1.69     | 3.2     |
| Independent  | 0.4B  | 80M      | 1.49     | 2.9     | 1.8      | 3.6     | 1.29     | 2.56     | 1.66     | 3.22     | 1.46     | 2.84     | 1.51          | 2.99     | 1.53     | 3.01     |

We can observe that in this setting, "Ours N=0" gives the highest speedup while "Ours N=1" has better acceptance length.
The trend we observed on Llama3-8B-Instruct is **maintained** (we are better than "EAGLE-2" and "EAGLE-2" is better than "Independent").
Once again, we would like to highlight that the speedup is highly dependent on the hardware used and we believe the gap will be bigger on more recent CPUs (see answer 2FbK.Q1 for more details https://openreview.net/forum?id=Rz0kozh3LE&noteId=yGJpPvjANr).
</content>
</thread_object>

</thread>
<thread>
creation date:
1738765937887

thread_type:
public_comment

<thread_object>
creator: ~penghui_yang1

<content>
title:
Clarification Needed on Overlap with ICML 2024 Work

comment:
Thank you for presenting this work. Upon reviewing the methodology described, I noticed striking similarities with the ICML 2024 paper, *GliDe with a CaPE: A Low-Hassle Method to Accelerate Speculative Decoding*. Both works appear to share a key process: draft tokens are processed through the large model's embedding layer, followed by self-attention to generate queries, which are then used to compute cross-attention with key-value pairs obtained from the large model.  

Given the strong resemblance, I believe it would greatly benefit the scientific discussion if this paper explicitly addressed the ICML 2024 work. A detailed comparison between the two would clarify the relationship and allow readers to better understand the unique contributions of this paper. Specifically, it would be helpful to highlight the novel aspects introduced here that are distinct from those in *GliDe with a CaPE*.  

Transparency and proper contextualization are critical for advancing the field, and I believe addressing this will further strengthen the paper. I look forward to seeing this clarification.
</content>
</thread_object>

</thread>
<thread>
creation date:
1740808252800

thread_type:
public_comment

<thread_object>
creator: ~penghui_yang1

<content>
title:
Misrepresentation of GliDe’s Methodology in the Camera-Ready Version

comment:
I would like to clarify a factual inaccuracy in the newly added comparison.  

In the updated version, the authors claim that *Du et al. (2024)* (i.e., GliDe) leverages the KV-cache of all layers of $M_{\text{Large}}$, implying that it requires as many layers in $M_{\text{Small}}$ as in $M_{\text{Large}}$ to fully capture the state. **This is incorrect.** GliDe explicitly states that it **only uses the KV-cache from the final layer of $M_{\text{Large}}$**, not all layers, and GliDe conducts an ablation study about this in Figure 7.  

This misrepresentation significantly affects the validity of the comparison. The argument that GliDe would result in slow drafting due to needing all layers is based on a false premise. I strongly encourage the authors to **correct this misunderstanding and provide a more accurate discussion** of GliDe’s approach.  

While the authors mention plans for an empirical comparison, at this stage, the two approaches appear **highly similar**, and I find it difficult to identify clear distinctions between them. I would appreciate it if the authors could further elaborate on the key differences to provide a clearer understanding.  

Looking forward to the authors’ clarification.
</content>
</thread_object>

</thread>


Metareview:
metareview:
All reviewers agreed that this is an good contribution to speculative decoding.
Strength:
1. An in-depth analysis of why new methodologies could be needed to improve upon the existing SD solutions.
2. An novel solution to mitigate the issues discussed.

Weakness:
1. Improvement over SOTA is not big. Given that the methodology is much more complex, so the potential impact can be small.