Kangaroo: Decodifica Autospeculativa Senza Perdite tramite Doppio Early Exiting

Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting

April 29, 2024

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

cs.AI

Abstract

La decodifica speculativa ha dimostrato la sua efficacia nell'accelerare l'inferenza dei modelli linguistici di grandi dimensioni mantenendo una distribuzione di campionamento coerente. Tuttavia, l'approccio convenzionale di addestrare un modello di bozza separato per ottenere un tasso di accettazione dei token soddisfacente può risultare costoso. Traendo ispirazione dall'early exiting, proponiamo un nuovo framework di decodifica auto-speculativa chiamato Kangaroo, che utilizza una sottorete fissa e poco profonda come modello di bozza auto-generato, con i rimanenti strati che fungono da modello target più grande. Addestriamo un modulo adattatore leggero ed efficiente sulla sottorete per colmare il divario tra la capacità rappresentativa della sottorete e quella del modello completo. È importante notare che la latenza di inferenza del modello di bozza auto-generato potrebbe non essere più trascurabile rispetto al modello grande, rendendo necessarie strategie per aumentare il tasso di accettazione dei token riducendo al minimo i passaggi di bozza del modello piccolo. Per affrontare questa sfida, introduciamo un ulteriore meccanismo di early exit per generare i token di bozza. Nello specifico, interrompiamo la previsione successiva del modello piccolo durante la fase di bozza una volta che il livello di confidenza per il token corrente scende al di sotto di una certa soglia. Esperimenti estensivi su Spec-Bench dimostrano l'efficacia di Kangaroo. Sotto verifica a sequenza singola, Kangaroo raggiunge accelerazioni fino a 1,68 volte su Spec-Bench, superando Medusa-1 con l'88,7% in meno di parametri aggiuntivi (67M rispetto a 591M). Il codice di Kangaroo è disponibile all'indirizzo 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.