KOALA: Enhancing Speculative Decoding for LLM via Multi-Layer Draft Heads with Adversarial Learning

要約

大規模言語モデル (LLM) は、自己回帰デコードの性質により、高い推論遅延を示します。
投機的デコードのドラフト ヘッドはこの問題を軽減しますが、その可能性の全容はまだ解明されていません。
この論文では、ドラフトヘッドへの直交アプローチである KOALA (K-layer Optimized Adversarial Learning Architecture) を紹介します。
従来の単層ドラフトヘッドを多層アーキテクチャに変換し、従来の教師ありトレーニングに敵対的学習を組み込むことにより、KOALA は後続のトークンを予測する際のドラフトヘッドの精度を大幅に向上させ、LLM の機能をより厳密に反映します。
この改善には、ドラフトのオーバーヘッドがわずかに増加するという代償が伴いますが、KOALA はドラフト ヘッドの可能性を大幅に解放し、投機的なデコードを大幅に強化します。
さまざまなタスクにわたる自己回帰ドラフト ヘッドと非自己回帰ドラフト ヘッドの両方を含む KOALA の包括的な評価を実施し、レイテンシーの高速化率が 0.24 倍から 0.41 倍向上したことを実証しました。これは、元のドラフト ヘッドよりも 10.57% ～ 14.09% 高速です。

要約(オリジナル)

Large Language Models (LLMs) exhibit high inference latency due to their autoregressive decoding nature. While the draft head in speculative decoding mitigates this issue, its full potential remains unexplored. In this paper, we introduce KOALA (K-layer Optimized Adversarial Learning Architecture), an orthogonal approach to the draft head. By transforming the conventional single-layer draft head into a multi-layer architecture and incorporating adversarial learning into the traditional supervised training, KOALA significantly improves the accuracy of the draft head in predicting subsequent tokens, thus more closely mirroring the functionality of LLMs. Although this improvement comes at the cost of slightly increased drafting overhead, KOALA substantially unlocks the draft head’s potential, greatly enhancing speculative decoding. We conducted comprehensive evaluations of KOALA, including both autoregressive and non-autoregressive draft heads across various tasks, demonstrating a latency speedup ratio improvement of 0.24x-0.41x, which is 10.57%-14.09% faster than the original draft heads.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google