Elamur Research
☰
home / papers / 2024 / EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive license to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/2401.15077

Summary

Generating text token by token makes LLM inference slow. Speculative sampling speeds it up by letting a cheap model propose several tokens that the full model then checks in a single pass. However, good draft models are hard to find, especially for the smallest model in a family. EAGLE replaces the separate draft model with a lightweight head: one transformer decoder layer that works on the target LLM's second-to-top-layer features, which follow more regular patterns than tokens. Predicting the next feature is ambiguous because it depends on which token was actually sampled. EAGLE resolves this by also feeding the head the token sequence shifted one step ahead. Draft tokens are produced through the original LM head and verified with standard speculative sampling, so the output distribution provably stays the same under greedy and non-greedy decoding. Tested on Vicuna, LLaMA2-Chat and Mixtral 8x7B Instruct across four benchmarks, EAGLE beats Lookahead and Medusa. It is cheap to train and needs no fine-tuning of the backbone.

Key points

  • On MT-bench, EAGLE runs 2.1x–3.8x faster than vanilla decoding, 1.7x–2.1x faster than Lookahead and 1.5x–1.6x faster than Medusa. For LLaMA2-Chat 70B it gives a 2.7x–3.5x latency speedup and roughly twice the throughput.
  • An ablation on Vicuna 7B shows how each design choice contributes:
    • Token-level drafting gives about a 1.5x speedup.
    • Feature-level drafting raises this to about 1.9x.
    • Adding the shifted tokens raises it to about 2.8x.
  • Draft accuracy is around 0.8, compared with roughly 0.6 for Medusa.
  • Training uses only ShareGPT data (at most 70k dialogues, about 2–4B tokens), with no training on the evaluation sets:
    • For the 70B model, the head has under 1B parameters and trains in 1–2 days on 4x A100 (40G).
    • Heads for the 7B–33B models can be trained on an RTX 3090 node.
  • Evaluation covers MT-bench, HumanEval, GSM8K and Alpaca. Combined with gpt-fast, EAGLE reaches 160.4 tokens/s for LLaMA2-Chat 7B on a single RTX 3090.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Computing & GPUs

Cited by

197 papers of our library cite it (of the 140673 whose bibliographies we read); 1 of them have a page here, 196 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2401.15077 · arXiv 2401.15077

@misc{li2024eagle,
  title = {{EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty}},
  author = {Li, Yuhui and Wei, Fangyun and Zhang, Chao and Zhang, Hongyang},
  year = {2024},
  doi = {10.48550/arxiv.2401.15077},
  eprint = {2401.15077},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2401.15077}
}

References

41 references, 5 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Cai, T., Li, Y., Geng, Z., Peng, H., and Dao, T. Medusa: Simple framework for accelerating LLM generation with multiple decoding heads. https://github.com/FasterDecoding/Medusa, 2023. ↗
  2. Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023a. find
  3. Chen, M., Tworek, J., Jun, H., Yuan, Q., Pinto, H. P. d. O., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021. find our card
  4. Chen, Z., Yang, X., Lin, J., Sun, C., Huang, J., and Chang, K. C.-C. Cascade speculative drafting for even faster LLM inference. arXiv preprint arXiv:2312.11462, 2023b. — [arXiv:2312.11462](https://arxiv.org/abs/2312.11462) find
  5. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021. find our card
  6. Fu, Y., Bailis, P., Stoica, I., and Zhang, H. Breaking the sequential dependency of LLM inference using lookahead decoding, November 2023. URL https://lmsys.org/blog/2023-11-21-lookahead-decoding/. ↗
  7. Gale, T., Elsen, E., and Hooker, S. The state of sparsity in deep neural networks.(2019). arXiv preprint cs.LG/1902.09574, 2019. find
  8. Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. find
  9. He, Z., Zhong, Z., Cai, T., Lee, J. D., and He, D. Rest: Retrieval-based speculative decoding. arXiv preprint arXiv:2311.08252, 2023. find
  10. Hinton, G., Vinyals, O., and Dean, J. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015. find our card
  11. Hooper, C., Kim, S., Mohammadzadeh, H., Genc, H., Keutzer, K., Gholami, A., and Shao, S. Speed: Speculative pipelined execution for efficient decoding. arXiv preprint arXiv:2310.12072, 2023. — [arXiv:2310.12072](https://arxiv.org/abs/2310.12072) find
  12. Hubara, I., Courbariaux, M., Soudry, D., El-Yaniv, R., and Bengio, Y. Quantized neural networks: Training neural networks with low precision weights and activations. journal of machine learning research, 18(187):1–30, 2018. find
  13. Jain, N., Chiang, P.-y., Wen, Y., Kirchenbauer, J., Chu, H.-M., Somepalli, G., Bartoldson, B. R., Kailkhura, B., Schwarzschild, A., Saha, A., et al. NEFTune: Noisy embeddings improve instruction finetuning. arXiv preprint arXiv:2310.05914, 2023. — [arXiv:2310.05914](https://arxiv.org/abs/2310.05914) find
  14. Kim, S., Gholami, A., Yao, Z., Mahoney, M. W., and Keutzer, K. I-bert: Integer-only bert quantization. In International conference on machine learning, pp. 5506– 5518. PMLR, 2021. find
  15. Kim, S., Mangalam, K., Moon, S., Malik, J., Mahoney, M. W., Gholami, A., and Keutzer, K. Speculative decoding with big little decoder. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. find
  16. Kurtic, E., Campos, D., Nguyen, T., Frantar, E., Kurtz, M., Fineran, B., Goin, M., and Alistarh, D. The optimal bert surgeon: Scalable and accurate second-order pruning for large language models. arXiv preprint arXiv:2203.07259, 2022. find
  17. Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274– 19286. PMLR, 2023. find
  18. Liu, X., Hu, L., Bailis, P., Stoica, I., Deng, Z., Cheung, A., and Zhang, H. Online speculative decoding. arXiv preprint arXiv:2310.07177, 2023. find
  19. Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Wong, R. Y. Y., Chen, Z., Arfeen, D., Abhyankar, R., and Jia, Z. SpecInfer: Accelerating generative LLM serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023. find
  20. Patterson, D. A. Latency lags bandwith. Communications of the ACM, 47(10):71–75, 2004. find
  21. PyTorch Labs. gpt-fast. https://github.com/pytorch-labs/gpt-fast/, 2023. ↗
  22. Sanh, V., Wolf, T., and Rush, A. Movement pruning: Adaptive sparsity by fine-tuning. Advances in Neural Information Processing Systems, 33:20378–20389, 2020. find
  23. Santilli, A., Severino, S., Postolache, E., Maiorca, V., Mancusi, M., Marin, R., and Rodola, E. Accelerating transformer inference for translation via parallel decoding. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Proceedings of the 61st Annual Meeting of the Association for Computationa… find
  24. Shazeer, N. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019. find
  25. Shen, S., Dong, Z., Ye, J., Ma, L., Yao, Z., Gholami, A., Mahoney, M. W., and Keutzer, K. Q-bert: Hessian based ultra low precision quantization of bert. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 8815–8821, 2020. find
  26. Spector, B. and Re, C. Accelerating LLM inference with staged speculative decoding. arXiv preprint arXiv:2308.04623, 2023. — [arXiv:2308.04623](https://arxiv.org/abs/2308.04623) find
  27. Stern, M., Shazeer, N., and Uszkoreit, J. Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31, 2018. find
  28. Sun, X., Ge, T., Wei, F., and Wang, H. Instantaneous grammatical error correction with shallow aggressive decoding. arXiv preprint arXiv:2106.04970, 2021. find
  29. Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang, P., and Hashimoto, T. B. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford\_alpaca, 2023. ↗
  30. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. LlAMA 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. find
  31. Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.09418, 2019. find our card
  32. Wu, Z., Liu, Z., Lin, J., Lin, Y., and Han, S. Lite transformer with long-short range attention. arXiv preprint arXiv:2004.11886, 2020. find
  33. Xia, H., Ge, T., Wang, P., Chen, S.-Q., Wei, F., and Sui, Z. Speculative decoding: Exploiting speculative execution for accelerating seq2seq generation. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 3909–3925, 2023. find
  34. Yang, N., Ge, T., Wang, L., Jiao, B., Jiang, D., Yang, L., Majumder, R., and Wei, F. Inference with reference: Lossless acceleration of large language models. arXiv preprint arXiv:2304.04487, 2023a. — [arXiv:2304.04487](https://arxiv.org/abs/2304.04487) find
  35. Yang, S., Lee, G., Cho, J., Papailiopoulos, D., and Lee, K. Predictive pipelined decoding: A compute-latency trade-off for exact llm decoding. arXiv preprint arXiv:2307.05908, 2023b. — [arXiv:2307.05908](https://arxiv.org/abs/2307.05908) find
  36. Zadeh, A. H., Edo, I., Awad, O. M., and Moshovos, A. Gobo: Quantizing attention-based nlp models for low latency and energy efficient inference. In 2020 53rd Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), pp. 811–824. find
  37. Zafrir, O., Boudoukh, G., Izsak, P., and Wasserblat, M. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing-NeurIPS Edition (EMC2-NIPS), pp. 36–39. find
  38. Zhang, J., Wang, J., Li, H., Shou, L., Chen, K., Chen, G., and Mehrotra, S. Draft & verify: Lossless large language model acceleration via self-speculative decoding. arXiv preprint arXiv:2309.08168, 2023. find
  39. Zhang, P., Zeng, G., Wang, T., and Lu, W. TinyLlama: An open-source small language model. arXiv preprint arXiv:2401.02385, 2024. find our card
  40. Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E., et al. Judging llm-as-a-judge with mt-bench and chatbot arena. arXiv preprint arXiv:2306.05685, 2023. find
  41. Zhou, Y., Lyu, K., Rawat, A. S., Menon, A. K., Rostamizadeh, A., Kumar, S., Kagy, J.-F., and Agarwal, R. DistillSpec: Improving speculative decoding via knowledge distillation. arXiv preprint arXiv:2310.08461, 2023. — [arXiv:2310.08461](https://arxiv.org/abs/2310.08461) find