Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding

Summary

This paper works on improving the drafting accuracy of speculative decoding. Authors mention a important problem that in the Medusa like speculative decoding methods, the draft heads make predictions without considering earlier tokens in the current candidate continuation. To solve this problem, they propose a drop-in sequentially dependent method to standard draft heads. This method increases the average candidate continuation acceptance length and further accelerates LLM inference.

Rating

7

Confidence

3

Ethics flag

1

Reasons to accept

1. The motivation is clear. The method is consistent with the motivation. The sequentially dependent method is novel and naturally considers earlier tokens in the current candidate continuation. 2. The experiments demonstrate that the proposed method significantly outperforms the baselines (2.70× acceleration). 3. Authors explore a lot on the architecture and training objective (section 3.1) . I especially appreciate the detailed analysis and discussions on this method.

Reasons to reject

The paper mainly compares with Medusa but lacks comparison with other speculative decoding methods. [1][2] [1] Li, Yuhui, et al. "Eagle: Speculative sampling requires rethinking feature uncertainty." arXiv preprint arXiv:2401.15077 (2024). [2] Chen, Charlie, et al. "Accelerating large language model decoding with speculative sampling." arXiv preprint arXiv:2302.01318 (2023).

Questions to authors

A suggestion: The task influences the speculative decoding method a lot, but the paper only evaluates on several benchmarks. According to this leaderboard (https://github.com/hemingkx/Spec-Bench), Hydra achieves a good performance on a wide range of tasks. It would be beneficial if the authors could provide additional experimental results to demonstrate the effectiveness of the method across different tasks.

Summary

This paper proposed a new speculative decoding scheme based on Medusa. This is different from Medusa which only considers the LLM's hidden state as the input and predicts the future tokens, this paper proposed to take the hidden stated and predicted candidates as input, which enhances the "dependencies" rather than fully factorized. Additionally, authors explored coupled improved model architecture and the training objectives, leading to the "Hydra++". Empirical results showed a two-digit margin increase w.r.t baselines regarding averaged accepted drafts, and throughput. In general, I think this is an elegant idea and tailored improvements to the model and objectives showed consistent gains.

Rating

8

Confidence

4

Ethics flag

1

Reasons to accept

- The idea itself is well-motivated and somewhat elegant to me. The predicted tokens are at hand information and could be used. It feels like adding an "autoregressiveness" in the draft heads. I would argue the idea is simple yet effective. - Authors also explored the design space of the model architectures, sampling methods, and training objectives. The design choices are well illustrated and motivated. And the final "hydra++" is pretty well-tuned. - The evaluation is comprehensive and empirical results are showing consistent gains even w.r.t. strong baselines. Additionally, the paper is pretty easy to follow and mathematical explanations are clear. In general, this paper is clear and solid with well-justified motivation and clean techniques. Empirical results also demonstrated gains.

Reasons to reject

- Sample methods: I would encourage authors to consider more sampling methods, for example, top-p. - Related work might consider adding some concurrent work?

Questions to authors

- Will the different sampling methods affect the gain over acceptance? - What are the other benchmarks that might be a good fit instead of MT-bench?

Summary

The paper introduces Hydra heads to introduce sequential dependency for the draft heads based on Medusa decoding. The paper conducts experiments using the Vicuna family of models to evaluate the throughput and token acceptance length compared to autoregressive decoding and Medusa and shows that the proposed method can improve on both metrics. The paper also investigates other aspects of the proposed Hydra methods by incorporating some recent techniques such as teacher loss and PrefixMLP to improve the framework further and thus presents Hydra++, which shows better performance with the combination of these techniques.

Rating

7

Confidence

3

Ethics flag

1

Reasons to accept

1. The presentation is clear and easy to understand. The motivation for incorporating sequential dependence to Medusa decoding is simple but seems effective from the experiments. 2. The studied problem is important and impactful for current LLM practices for decoding efficiencies. 3. The exploration of other techniques is quite extensive seeking alternative implementations of Hydra heads. The results and observations can provide useful insights for practitioners.

Reasons to reject

1. It would be good to have a cost and efficiency analysis for the proposed method since taking the previous tokens to the draft heads seems to scale with the number of heads. 2. The optimal draft head recipe section is not comprehensive enough to understand the intuitions behind some observations. For example, could the authors elaborate more on why adding noise to Hydra did not improve the performance of the models? The motivation to add noise here is unclear since the referenced paper is about using noisy embeddings for instruction tuning. It would be better to see if adding noise to the speculative decoding generally fails or if there are other observations for the relationship between noisy embeddings and decoding performance. 3. Other than greedy decoding, Some results on other temperatures for evaluation would be interesting to see how randomness would affect the observations.

Questions to authors

1. See the weakness above. 2. For Figures 2 and 3, why is there a confidence interval plotted (if I understand correctly) for the Hydra only on Avg. Throughput results? Please correct me if I miss anything. 3. Typo: the word "using" is repeated twice in the 2nd line of section 7.

Thank you for the response

Thanks to the authors for their responses. I think the new results and explanations are reasonable and expected. I would like to maintain my current evaluations.