Mixture of Attentions For Speculative Decoding

Thank you very much for your feedback and suggestions. We are delighted that you found the discussion, analysis and evaluation, particularly the client-server use case, valuable and interesting. Let us address your suggestions that will help us improve the paper further.

Q1: 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.

The primary reason is that, with Layer Self-Attention, we can observe the whole dynamical system state whereas EAGLE or Medusa would only observe a limited (transformed) subset of this state.

Q2: 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.

Thank you for raising this very interesting question. Although we focused on improving the current SOTA method (EAGLE-2), our observations (partial observability, on-policyness and target inference layer) are true for many self-drafting methods, for instance, it could also be applied to MEDUSA, MLP Speculator (https://arxiv.org/abs/2404.19124) or [4, 5]. We will make this clearer in the final version of the paper.

Regarding non-self-drafting SD, this should be studied on a case-by-case basis. For instance, target inference layer could potentially be applied to independent small models. Many student-teacher distillation frameworks [1-2], already leverage the on-policyness property by generating directly from the student but are mostly are 1-step bounded (therefore expensive to train). For SD methods based on lookahead decoding, it would generally not apply. One exception is Ouroboros [3] that leverages a small model with lookahead decoding. Their small model could also benefit from our solutions.

[1] Gu, Yuxian, et al. Knowledge distillation of large language models. ICLR, 2024.

[2] Zhou, YYongchao, et al. Distillspec: Improving speculative decoding via knowledge distillation. ICLR 2024.

[3] Zhao, Weilin, et al. "Ouroboros: Speculative Decoding with Large Model Enhanced Drafting." EMNLP 2024.

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

[5] Kim et al. (2024), Accelerating Blockwise Parallel Language Models with Draft Refinement.

Q3: 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.

In Table 3, we show (accuracy-like) pass@k scores on HumanEval in a client-server setting, as an example of performance after a partial/full disconnection. Tables 2 and 3 are focused on efficiency, citing standard SD metrics such as acceptance length and speedup ratios. This practice is very common for SD papers as the verification step of SD is theoretically shown to preserve the output distribution of the large LLM.

The difference you observed between SD and non-SD methods is probably due to how verification is implemented. Therefore, mitigating the partial observability will not overcome this problem as it does not touch the verification algorithm. If you observed a difference in greedy decoding, it is likely due to the (b)float16 precision. Depending on the attention implementation and the precision used, doing a forward pass with a batched request of N tokens (what is done during verification of SD methods) and doing N forward pass of 1 token (vanilla decoding) can lead to very tiny differences in logits.

• Q1. Besides intuition of each part, I wonder if the author can provide an overall review of each component, or if they are implemented as a bag of tricks
• Q2. Since the work only compare with EAGLE, it is important to establish the single reference point (with clear arguments) not only against leading SD methods but also other alternatives (such as lora, quantization, or simply training smaller lm better)
• Q4: with same concern as Q2, I agree with Reviewer Sr5b that general performance metrics such as accuracy or users study should be included to assess the validity of the final llm output. Only providing SD metrics leads to the concern that the improvement is only made by fooling the validating model instead of genuine improvement.

Thank you very much for your valuable feedback. We are delighted that you found our background on SD informative and the client-server evaluation interesting and useful. Let us address the areas for improvement that will make the paper submission stronger.

Q1: 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.

We will revise the paper to enhance clarity. Here are some additional intuitions of our proposed components.

• Partial Observability: EAGLE or Medusa only observe a limited (transformed) subset of the whole dynamical state of the LLM. With Layer Self-Attention we can observe the whole (non transformed) state.
• On-policyness ensures that the small model is trained under conditions that closely resemble the actual inference process, i.e. the LLM state is not observed at every time-step. However, training fully on-policy is expensive because it requires generation during training. With our Cross-Attention solution, we can train more on-policy without having to generate because we can hide specific keys and values.
• Target Layer Inference: we can reuse the last N layers of the LLM to simplify the inference of the small model by predicting the input of the Nth layer instead of the final layer.

Q2: 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.

EAGLE-2 evaluated a wide range of LLMs from 7B to 70B parameters, which show predictable and consistent improvements across model scales and types. They also compare against (and improve upon) other methods such as Medusa, Lookahead and Speculative Sampling, establishing themselves as the leading SD methodology [1][2]. Unfortunately, we do not have the same compute capacity as them thus choose one of EAGLE's SOTA LLMs, Llama3, for a fair comparison. We will update Section 4 appropriately, thank you for pointing this out.

[1] Li, Yuhui et al. “EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees.” Conference on Empirical Methods in Natural Language Processing (2024). https://arxiv.org/pdf/2406.16858

[2] Xia, Heming et al. “Unlocking Efficiency in Large Language Model Inference: A Comprehensive Survey of Speculative Decoding.” Annual Meeting of the Association for Computational Linguistics (2024).

Q3: the hyperparameter N is used many times. Giving it a name can help readability.

We introduce the concept of "Target Layer Inference" in Section 3.2 hence we could replace "N" with "Target Layer", or just "TL" for short. Thanks for the suggestion.

Q4: 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.

Although it is technically possible that the small model generations could be better than what the large model would generate (in terms of accuracy, safety or other non-speculative metrics), being better than the large model is usually not within the scope of speculative sampling because the small models are assumed to be less capable.

(To be continued ...)

Thank you so much for your valuable feedback. We are really happy that you appreciated the clarity and motivation of the paper as well as our additional client-server experiments aimed to evaluate the method in a very realistic use case. Your enthusiasm for our work is really great so let us address your suggestions for further improvement.

Q1: 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.

In EAGLE [1] and Medusa [2], the authors evaluate a wide range of LLMs from 7B to 70B parameters, which show predictable and consistent improvements across model scales and types. Unfortunately, we do not have the same compute capacity as them thus choose one of EAGLE-2's SOTA LLMs, Llama3, for a fair comparison. Nevertheless, we are currently training small models for Qwen2.5-3B-Instruct to confirm that we will obtain similar results. We hope we can obtain the results before the end of the discussion period.

We note that our model is always better than EAGLE-2 for every dataset (with N=1). We would also like to emphasise that the speed up ratios are highly dependent on the hardware used (notably the CPU because of the complexity of the EAGLE-2 tree-decoding strategy). For instance, it was impossible for us to reproduce the 3.46x speedup claimed by EAGLE-2 on our hardware (we achieved 2.14x for the exact same setup). This has been reported by others (see https://github.com/SafeAILab/EAGLE/issues/121). Therefore, we believe that on more recent CPU, the gap with EAGLE-2 and our architecture would be more significant as we have much better acceptance length.

[1] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang. International Conference on Machine Learning, 2024.

[2] Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, Tri Dao. International Conference on Machine Learning, 2024.

Q2: 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. [...]

Apologies for any misunderstanding. We believe the client-server evaluation is a good example of an application where our method provides a clear benefit. For example, we beat EAGLE-2 on decoding speeds (up to 80% on 4G and up to 40% on 5G) while lowering the calls to the Large LLM by around 50%. The contribution goes beyond efficiency gains, for example, in the event of a total disconnection, our approach maintains high accuracy compared to other SD methods (15% - 38% higher pass@1 score on HumanEval) and demonstrates advantages over API calls to LLMs, which would otherwise be unable to continue the generation process. One example could be a user working on a laptop with an LLM integrated system or browser on a train where its internet connection is not stable. Another example could be LLM-based translation on mobile phones, etc. We will update the introduction accordingly to highlight more concrete use cases.

Regarding the scope of the contributions, although we focused on improving the current SOTA method (EAGLE-2), our observations (partial observability, on-policyness, target inference layer) are true for many self-drafting methods, for instance, MEDUSA or MLP Speculator (https://arxiv.org/abs/2404.19124) should also benefit from those. Therefore, we believe that self-drafting architectures should be aware of our proposals.

Q3: 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?

We chose Llama3 (8B) as it is a state-of-the-art LLM and because we could compare it to the EAGLE-2 checkpoint. Moreover, there is no fundamental difference between decoder-only transformers architectures that would require significant adaptations to our approach. As previous works showed, the performance of self-drafting methods is consistent across models and model sizes, see answer to Q1. We will update the paper to reflect this.

Thank you for your thorough review and constructive feedback. We appreciate your insights and will address each of the points raised to improve the quality of our paper. We are pleased that you acknowledge the importance of our contributions, particularly the client-server framework and the introduction of Layer Self-Attention (LSA) and Cross-Attention (CA) layers to address partial observability and on-policyness. We believe those are important for the future development of self-drafting small models.

Q1. 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.

We agree that the number of parameters and the training data are crucial. That is exactly why we ran our benchmarking in a very controlled setting. We have trained the relevant baselines (EAGLE-2, Independent small models) and our proposed solutions using the same training dataset (Ultrachat 200k) and number of epochs for a fair comparison. For reference, we have included the performance of the official EAGLE-2 checkpoint that was trained on the ShareGPT data. We did not use ShareGPT as it was unclear which version of the dataset was used.

We also took great care to match the number of model parameters, i.e. "Ours (N=0)", "EAGLE-2", "EAGLE-2 off." and "Independent 1.3B" all have 1.3B parameters (250M trainable and 1.05B frozen, for the "LM head" and "token embed" layers). We chose 250M trainable parameters to be directly comparable to EAGLE-2 and their official checkpoint. Since increasing N also increases the total number of parameters (frozen in this case), we added a larger "Independent 1.7B" model to compare it with "Ours (N=1)" using 1.55B parameters and "Ours (N=3)" with 1.8B parameters. Table 2 clearly shows an advantage for us in both cases.

Finally, all decoding methods use the same tree-decoding strategy. We will clarify this in the updated version of the paper and include the reference to Yi et al. (2024).

Q2. Discussion for the memory-bound nature of LLM is required in the paper.

Indeed, we will add a paragraph in the background section to explain that for smaller batch sizes, the LLM inference is memory-bound and that speculative decoding can better leverage the spare compute especially with tree-attention.

Q3. 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.

Apologies for any misunderstanding, however, Llama3 (8B) is a SOTA LLM between 3B and 13B. In EAGLE [1] and Medusa [2], the authors evaluate a wide range of LLMs from 7B to 70B parameters, which show predictable and consistent improvements across model scales and types. Unfortunately, we do not have the same compute capacity as them thus choose one of EAGLE-2's SOTA LLMs, Llama3, for a fair comparison. Moreover, our small models are bounded by the practicalities of small device memory for the client-server setting. Nevertheless, we are currently training small models for Qwen2.5-3B-Instruct to confirm that we will obtain similar results. We hope we can obtain the results before the end of the discussion period.

[1] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang. International Conference on Machine Learning, 2024.

[2] Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, Tri Dao. International Conference on Machine Learning, 2024.

Q4. Typo? line 466.

Sorry, could you please quote the exact string to help us pinpoint the possible typo? Thank you.

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

Thank you so much for the citations. Rest assured we will add them to the discussions in related work.

We believe our proposed Layer Self Attention and Target Layer Inference could to easily be applied in [1]. Similarly, we could also mitigate partial observability in [2] with Layer Self Attention.

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

[2] Kim et al. (2024), Accelerating Blockwise Parallel Language Models with Draft Refinement.

Thank you for your thorough responses to the points I raised. Upon reviewing your clarifications and additional explanations, I realize that some of my earlier concerns were based on misunderstandings or insufficiently detailed observations. I appreciate the effort you have made to provide clear and specific information to address the weaknesses I highlighted. The practical advancements in client-server frameworks and the methodological contributions addressing partial observability and on-policyness are significant and relevant to the field. I raised the score, accordingly.

Thank you for sharing this paper; we were not aware of it. We have added it to the related work section and discussed the differences with it. We aim to include a comparison with this approach, unfortunately no trained checkpoints are available online. It will take some time to train it in the same scenario for a fair comparison.

Thank you for your message.

We did not write that Glide always needs the same number of layers in $M_\text{Small}$ as what is in $M_\text{Large}$. What we wrote is that, if Glide wanted to fully capture the state of the dynamical system, which you should do if you want to maximize acceptance rate [our first contribution is to identify this], then Glide will need the same number of layers inside $M_\text{Small}$ and $M_\text{Large}$ according to their architecture description.

Now, as in practice they only use the last layer KV-cache, they will suffer from partial observability (see our paper). We confirmed this by adding Glide in our experiments (Table 2). It is slower than EAGLE-2 but has higher acceptance rate. However, compared to us, their acceptance rate is indeed lower as we would expect from our theory.

So, how can we capture the state of the dynamical system without needing many layers in $M_\text{Small}$? With our proposed Layer Self Attention layer [second contribution/key difference].

Therefore, what we wrote in the discussion section is correct.

Regarding other distinctions with Glide, we also propose target layer inference, and we try to always justify our design choices with formal properties (for instance why is it important to have a self attention layer before the cross attention one). We also focused on client-server scenarios and released the source code of the 4 methods (Mixture of Attentions, EAGLE-2, Glide and Independent) within a single framework.