Speculative Decoding via Early-exiting for Faster LLM Inference with Thompson Sampling Control Mechanism

Jiahao Liu, Qifan Wang, Jingang Wang, Xunliang Cai


Abstract
The recent advancements in large language models (LLMs) have been extraordinary, yet the escalating inference costs associated with them present challenges in real-world applications. To address these challenges, we propose a novel approach called Early-exiting Speculative Decoding (EESD) with lossless acceleration. Specifically, EESD utilizes a segment of the LLM to generate draft tokens, incorporating Early-exiting structures after the first N layers. To enhance the quality of draft tokens, a self-distillation method is integrated. This early-exiting design not only reduces deployment and training costs but also significantly accelerates the token generation speed. Moreover, we introduce a novel sampling mechanism that leverages Thompson Sampling to regulate the generation processes, automatically determining the quantity of draft tokens in each round. The original LLM is then employed to validate these draft tokens through a single forward pass, and thus guarantees that the final output text maintains a distribution consistent with vanilla auto-regressive decoding. The experimental results on both 13B and 70B models demonstrate that our approach decodes tokens at a markedly accelerated rate compared to prior methods, showing the effectiveness of our approach.
Anthology ID:
2024.findings-acl.179
Volume:
Findings of the Association for Computational Linguistics: ACL 2024
Month:
August
Year:
2024
Address:
Bangkok, Thailand
Editors:
Lun-Wei Ku, Andre Martins, Vivek Srikumar
Venue:
Findings
SIG:
Publisher:
Association for Computational Linguistics
Note:
Pages:
3027–3043
Language:
URL:
https://aclanthology.org/2024.findings-acl.179
DOI:
10.18653/v1/2024.findings-acl.179
Bibkey:
Cite (ACL):
Jiahao Liu, Qifan Wang, Jingang Wang, and Xunliang Cai. 2024. Speculative Decoding via Early-exiting for Faster LLM Inference with Thompson Sampling Control Mechanism. In Findings of the Association for Computational Linguistics: ACL 2024, pages 3027–3043, Bangkok, Thailand. Association for Computational Linguistics.
Cite (Informal):
Speculative Decoding via Early-exiting for Faster LLM Inference with Thompson Sampling Control Mechanism (Liu et al., Findings 2024)
Copy Citation:
PDF:
https://aclanthology.org/2024.findings-acl.179.pdf