Kangaroo: Verliesvrije Zelf-Speculatieve Decodering via Dubbele Vroege Uitstap

Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting

April 29, 2024

Auteurs: Fangcheng Liu, Yehui Tang, Zhenhua Liu, Yunsheng Ni, Kai Han, Yunhe Wang

cs.AI

Samenvatting

Speculatieve decodering heeft zijn effectiviteit aangetoond bij het versnellen van de inferentie van grote taalmodelen, terwijl een consistente steekproefverdeling behouden blijft. De conventionele aanpak van het trainen van een apart draft-model om een bevredigend tokenacceptatiepercentage te bereiken, kan echter kostbaar zijn. Geïnspireerd door early exiting, stellen we een nieuw zelf-speculatief decoderingframework voor, genaamd Kangaroo, dat een vast ondiep sub-netwerk gebruikt als een zelf-draft-model, waarbij de overige lagen dienen als het grotere doelmodel. We trainen een lichtgewicht en efficiënt adaptermodule bovenop het sub-netwerk om de kloof tussen het sub-netwerk en de representatiecapaciteit van het volledige model te overbruggen. Het is opmerkelijk dat de inferentielatentie van het zelf-draft-model niet langer verwaarloosbaar is in vergelijking met het grote model, wat strategieën vereist om het tokenacceptatiepercentage te verhogen terwijl de draftstappen van het kleine model worden geminimaliseerd. Om deze uitdaging aan te pakken, introduceren we een extra early exiting-mechanisme voor het genereren van draft-tokens. Specifiek stoppen we de daaropvolgende voorspelling van het kleine model tijdens de draftfase zodra het betrouwbaarheidsniveau voor het huidige token onder een bepaalde drempelwaarde valt. Uitgebreide experimenten op de Spec-Bench tonen de effectiviteit van Kangaroo aan. Onder single-sequence verificatie behaalt Kangaroo versnellingen tot 1,68 keer op de Spec-Bench, wat beter presteert dan Medusa-1 met 88,7% minder extra parameters (67M vergeleken met 591M). De code voor Kangaroo is beschikbaar op https://github.com/Equationliu/Kangaroo.

English

Speculative decoding has demonstrated its effectiveness in accelerating the inference of large language models while maintaining a consistent sampling distribution. However, the conventional approach of training a separate draft model to achieve a satisfactory token acceptance rate can be costly. Drawing inspiration from early exiting, we propose a novel self-speculative decoding framework Kangaroo, which uses a fixed shallow sub-network as a self-draft model, with the remaining layers serving as the larger target model. We train a lightweight and efficient adapter module on top of the sub-network to bridge the gap between the sub-network and the full model's representation ability. It is noteworthy that the inference latency of the self-draft model may no longer be negligible compared to the large model, necessitating strategies to increase the token acceptance rate while minimizing the drafting steps of the small model. To address this challenge, we introduce an additional early exiting mechanism for generating draft tokens. Specifically, we halt the small model's subsequent prediction during the drafting phase once the confidence level for the current token falls below a certain threshold. Extensive experiments on the Spec-Bench demonstrate the effectiveness of Kangaroo. Under single-sequence verification, Kangaroo achieves speedups up to 1.68times on Spec-Bench, outperforming Medusa-1 with 88.7\% fewer additional parameters (67M compared to 591M). The code for Kangaroo is available at https://github.com/Equationliu/Kangaroo.