Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding

要約

自己回帰 LLM 推論のメモリ帯域幅制限の性質に対抗するために、以前の研究では投機的デコード フレームワークが提案されています。
投機的デコードを実行するために、小規模なドラフト モデルは入力シーケンスの継続候補を提案し、その後、基本モデルによって並行して検証されます。
最近の Medusa デコード フレームワークで使用されているように、ドラフト モデルを指定する 1 つの方法は、ベース モデルの隠れ状態で動作する、ドラフト ヘッドと呼ばれる軽量ヘッドのコレクションとして指定することです。
これまでのところ、既存のドラフト ヘッドはすべて順番に独立しています。つまり、候補継続内の先行するトークンとは無関係に、候補継続中のトークンを推測します。
この研究では、ドラフト ヘッドの推測の精度を大幅に向上させる、標準ドラフト ヘッドの逐次依存型ドロップイン代替品である Hydra ヘッドを提案します。
私たちは、ヒドラ ヘッドのトレーニング目標とアーキテクチャの設計空間をさらに調査し、慎重に調整されたヒドラ ヘッド レシピ (Hydra++ と呼ぶ) を提案します。これにより、メデューサ デコードと自己回帰デコードと比較して、デコード スループットがそれぞれ最大 1.31 倍と 2.70 倍向上します。
。
全体として、Hydra ヘッドは、標準のドラフト ヘッドに対するシンプルかつ十分な動機を持った介入であり、ドラフト ヘッドベースの投機的デコードのエンドツーエンドの速度を大幅に向上させます。
私たちはコードを https://github.com/zankner/Hydra で公開しています。

要約(オリジナル)

To combat the memory bandwidth-bound nature of autoregressive LLM inference, previous research has proposed the speculative decoding frame-work. To perform speculative decoding, a small draft model proposes candidate continuations of the input sequence that are then verified in parallel by the base model. One way to specify the draft model, as used in the recent Medusa decoding framework, is as a collection of lightweight heads, called draft heads, that operate on the base model’s hidden states. To date, all existing draft heads have been sequentially independent, meaning that they speculate tokens in the candidate continuation independently of any preceding tokens in the candidate continuation. In this work, we propose Hydra heads: a sequentially-dependent drop-in replacement for standard draft heads that significantly improves the accuracy of draft head speculation. We further explore the design space of Hydra head training objectives and architectures, and propose a carefully tuned Hydra head recipe, which we call Hydra++, that improves decoding throughput by up to 1.31x and 2.70x compared to Medusa decoding and autoregressive de-coding respectively. Overall, Hydra heads are a simple and well-motivated intervention on standard draft heads that significantly improve the end-to-end speed of draft head-based speculative decoding. We make our code publicly available at https://github.com/zankner/Hydra.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google