Känguru: Verlustfreies selbstspekulatives Dekodieren durch doppeltes frühzeitiges Beenden.

Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting

April 29, 2024

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

cs.AI

Zusammenfassung

Spekulatives Decodieren hat seine Wirksamkeit bei der Beschleunigung der Inferenz großer Sprachmodelle unter Beibehaltung einer konsistenten Stichprobenverteilung gezeigt. Allerdings kann der herkömmliche Ansatz, ein separates Entwurfsmodell zu trainieren, um eine zufriedenstellende Token-Akzeptanzrate zu erreichen, kostspielig sein. Inspiriert von frühzeitigem Abbruch schlagen wir ein neuartiges selbstspekulatives Decodierungs-Framework namens Kangaroo vor, das ein festes flaches Teilnetzwerk als Selbstentwurfsmodell verwendet, wobei die verbleibenden Schichten als das größere Zielmodell dienen. Wir trainieren ein leichtes und effizientes Adaptermodul oben auf dem Teilnetzwerk, um die Kluft zwischen der Fähigkeit des Teilnetzwerks und der Repräsentationsfähigkeit des vollständigen Modells zu überbrücken. Es ist erwähnenswert, dass die Inferenzlatenz des Selbstentwurfsmodells im Vergleich zum großen Modell möglicherweise nicht mehr vernachlässigbar ist, was Strategien erfordert, um die Token-Akzeptanzrate zu erhöhen und gleichzeitig die Entwurfschritte des kleinen Modells zu minimieren. Um diese Herausforderung anzugehen, führen wir einen zusätzlichen frühzeitigen Abbruchmechanismus für die Generierung von Entwurfstoken ein. Konkret unterbrechen wir die nachfolgende Vorhersage des kleinen Modells während der Entwurfsphase, sobald der Vertrauensgrad für das aktuelle Token unter einen bestimmten Schwellenwert fällt. Umfangreiche Experimente auf dem Spec-Bench zeigen die Wirksamkeit von Kangaroo. Unter Einzel-Sequenz-Verifikation erzielt Kangaroo auf dem Spec-Bench Geschwindigkeitssteigerungen von bis zu 1,68-mal und übertrifft Medusa-1 mit 88,7\% weniger zusätzlichen Parametern (67M im Vergleich zu 591M). Der Code für Kangaroo ist verfügbar unter 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.