Kangaroo : Décodage auto-spéculatif sans perte via une double sortie précoce

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

Résumé

Le décodage spéculatif a démontré son efficacité pour accélérer l'inférence des grands modèles de langage tout en maintenant une distribution d'échantillonnage cohérente. Cependant, l'approche conventionnelle consistant à entraîner un modèle de brouillon séparé pour atteindre un taux d'acceptation de jetons satisfaisant peut s'avérer coûteuse. S'inspirant de la sortie anticipée, nous proposons un nouveau cadre de décodage auto-spéculatif, Kangaroo, qui utilise un sous-réseau peu profond fixe comme modèle de brouillon auto-généré, les couches restantes servant de modèle cible plus large. Nous entraînons un module d'adaptation léger et efficace sur ce sous-réseau pour combler l'écart entre les capacités de représentation du sous-réseau et celles du modèle complet. Il est à noter que la latence d'inférence du modèle de brouillon auto-généré peut ne plus être négligeable par rapport au grand modèle, nécessitant des stratégies pour augmenter le taux d'acceptation des jetons tout en minimisant les étapes de génération du petit modèle. Pour relever ce défi, nous introduisons un mécanisme supplémentaire de sortie anticipée pour générer les jetons de brouillon. Plus précisément, nous interrompons la prédiction ultérieure du petit modèle pendant la phase de génération une fois que le niveau de confiance pour le jeton actuel tombe en dessous d'un certain seuil. Des expériences approfondies sur le Spec-Bench démontrent l'efficacité de Kangaroo. Sous vérification de séquence unique, Kangaroo atteint des accélérations allant jusqu'à 1,68 fois sur Spec-Bench, surpassant Medusa-1 avec 88,7 % de paramètres supplémentaires en moins (67M contre 591M). Le code de Kangaroo est disponible à l'adresse 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.