Optimal Block-Level Draft Verification for Accelerating Speculative Decoding

要約

投機的デコードは、推論中に大規模言語モデル (LLM) をロスレスで高速化する効果的な方法であることが示されています。
各反復では、アルゴリズムは最初に小さなモデルを使用してトークンのブロックを作成します。
その後、トークンは大規模モデルによって並行して検証され、最終出力が大規模モデルの分布に従っていることを保証するために、トークンのサブセットのみが保持されます。
これまでのすべての投機的デコード作業では、ドラフト検証はトークンごとに独立して実行されます。
この研究では、追加の計算コストやドラフト トークンを発生させることなく、実時間のさらなる高速化を実現する、より優れたドラフト検証アルゴリズムを提案します。
まず、ドラフト検証ステップをブロックレベルの最適トランスポート問題として定式化します。
ブロックレベルの定式化により、より広範囲のドラフト検証アルゴリズムを検討し、1 つのドラフト ブロックで予想されるより多くの受け入れられたトークンを取得できるようになります。
ブロックレベルのトランスポート問題に対して最適な許容長を達成する検証アルゴリズムを提案します。
私たちは、提案したブロックレベルの検証アルゴリズムを幅広いタスクやデータセットで経験的に評価し、トークンレベルの検証アルゴリズムと比較して実時間の高速化が一貫して向上していることを観察しました。
私たちの知る限り、私たちの研究は、より優れたドラフト検証アルゴリズムによる投機的デコードの改善を確立した最初のものです。

要約(オリジナル)

Speculative decoding has shown to be an effective method for lossless acceleration of large language models (LLMs) during inference. In each iteration, the algorithm first uses a smaller model to draft a block of tokens. The tokens are then verified by the large model in parallel and only a subset of tokens will be kept to guarantee that the final output follows the distribution of the large model. In all of the prior speculative decoding works, the draft verification is performed token-by-token independently. In this work, we propose a better draft verification algorithm that provides additional wall-clock speedup without incurring additional computation cost and draft tokens. We first formulate the draft verification step as a block-level optimal transport problem. The block-level formulation allows us to consider a wider range of draft verification algorithms and obtain a higher number of accepted tokens in expectation in one draft block. We propose a verification algorithm that achieves the optimal accepted length for the block-level transport problem. We empirically evaluate our proposed block-level verification algorithm in a wide range of tasks and datasets, and observe consistent improvements in wall-clock speedup when compared to token-level verification algorithm. To the best of our knowledge, our work is the first to establish improvement over speculative decoding through a better draft verification algorithm.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google