Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting

要約

投機的デコードは、一貫したサンプリング分布を維持しながら、大規模な言語モデルの推論を高速化する効果があることが実証されています。
ただし、満足のいくトークン受け入れ率を達成するために別のドラフト モデルをトレーニングする従来のアプローチでは、コストがかかる可能性があります。
初期の終了からインスピレーションを得て、私たちは新しい自己投機的復号フレームワーク \emph{Kangaroo} を提案します。これは、固定された浅いサブネットワークを自己ドラフト モデルとして使用し、残りの層がより大きなターゲット モデルとして機能します。
サブネットワークと完全なモデルの表現能力の間のギャップを埋めるために、サブネットワーク上で軽量で効率的なアダプター モジュールをトレーニングします。
セルフドラフトモデルの推論レイテンシーは大規模モデルと比較して無視できなくなる可能性があり、小規模モデルのドラフト手順を最小限に抑えながらトークン受け入れ率を高める戦略が必要になる可能性があることは注目に値します。
この課題に対処するために、ドラフト トークンを生成するための追加の早期終了メカニズムを導入します。
具体的には、現在のトークンの信頼レベルが特定のしきい値を下回ると、草案段階で小規模モデルの後続の予測を停止します。
Spec-Bench での広範な実験により、カンガルーの有効性が実証されました。
単一シーケンス検証では、Kangaroo は Spec-Bench で最大 $1.68\times$ の高速化を達成し、88.7\% 少ない追加パラメーター (591M と比較して 67M) で Medusa-1 を上回ります。
Kangaroo のコードは https://github.com/Equationliu/Kangaroo で入手できます。

要約(オリジナル)

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 \emph{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.68\times$ 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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google