[論文レビュー] SpecTr: Fast Speculative Decoding via Optimal Transport

Ziteng Sun, Ananda Theertha Suresh+4 more|arXiv (Cornell University)|Oct 23, 2023

Natural Language Processing Techniques被引用数 4

ひとこと要約

SpecTr は、最適輸送理論を用いて高品質なドラフトトークンを選択する、画期的な推論手法を提案する。これにより、品質劣化を伴わず、自己回帰的生成が高速化される。この手法は、ベースライン自己回帰的デコードに対して 2.13 倍のウォールクロックスピードアップを達成し、標準的な推論手法よりも 1.37 倍の高速化を実現した。

ABSTRACT

Autoregressive sampling from large language models has led to state-of-the-art results in several natural language tasks. However, autoregressive sampling generates tokens one at a time making it slow, and even prohibitive in certain tasks. One way to speed up sampling is $ extit{speculative decoding}$: use a small model to sample a $ extit{draft}$ (block or sequence of tokens), and then score all tokens in the draft by the large language model in parallel. A subset of the tokens in the draft are accepted (and the rest rejected) based on a statistical method to guarantee that the final output follows the distribution of the large model. In this work, we provide a principled understanding of speculative decoding through the lens of optimal transport (OT) with $ extit{membership cost}$. This framework can be viewed as an extension of the well-known $ extit{maximal-coupling}$ problem. This new formulation enables us to generalize the speculative decoding method to allow for a set of $k$ candidates at the token-level, which leads to an improved optimal membership cost. We show that the optimal draft selection algorithm (transport plan) can be computed via linear programming, whose best-known runtime is exponential in $k$. We then propose a valid draft selection algorithm whose acceptance probability is $(1-1/e)$-optimal multiplicatively. Moreover, it can be computed in time almost linear with size of domain of a single token. Using this $new draft selection$ algorithm, we develop a new autoregressive sampling algorithm called $ extit{SpecTr}$, which provides speedup in decoding while ensuring that there is no quality degradation in the decoded output. We experimentally demonstrate that for state-of-the-art large language models, the proposed approach achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.

研究の動機と目的

• メンバーーシップコストを用いた最適輸送に基づく、推論のための原理的枠組みを構築すること。
• kトークンのドラフトをトークンレベルで一般化し、ドラフト品質と受容率を向上させること。
• 1−1/e の近似最適性を達成するドラフト選択アルゴリズムを設計し、1トークンあたりほぼ線形時間で実行可能とすること。
• 最終出力分布が大規模言語モデルの真の分布と一致することを保証し、品質劣化がないことを保証すること。
• 最先端のモデルを用いた実験的検証を行い、実世界の推論において顕著な高速化を示すこと。

提案手法

• メンバーーシップコストを用いた最適輸送問題として推論を定式化し、最大カップリングフレームワークを拡張する。
• 受容確率を最大化するように輸送計画を最適化する kトークンドラフト選択戦略を導入する。
• 線形計画法により最適なドラフト選択を解くが、1−1/e の近似最適性を達成する実用的近似を提案する。
• 1トークンあたりの語彙サイズにほぼ線形な時間でドラフト選択を計算する高速アルゴリズムを設計する。
• 大規模モデルがドラフトトークンを並列にスコアリングできる新しい自己回帰的サンプリングアルゴリズム SpecTr にドラフト選択を統合する。
• 統計的受容メカニズムを用いて、最終出力が大規模モデルの真の分布に従うことを保証する。

実験結果

リサーチクエスチョン

• RQ1 最適輸送理論を用いて、推論のための原理的枠組みを提供できるか？
• RQ2 推論を kトークンドラフトをトークンレベルで一般化することで、受容率を向上させられるか？
• RQ3 最適輸送定式化下でのドラフト選択品質の理論的限界は何か？
• RQ4 1トークンあたりほぼ線形時間で、近似的に最適なドラフト選択アルゴリズムを効率的に計算できるか？
• RQ5 提案手法は、出力品質を劣化させることなく、実際の応用で顕著な高速化を達成できるか？

主な発見

• LM1Bベンチマークにおいて、大規模モデルに PALM-2-Bison を用いた場合、SpecTr はベースライン自己回帰的デコードに対して 2.13 倍のウォールクロックスピードアップを達成した。
• L=8 かつ K=8 の条件下で、SpecTr は 4.0 のブロック効率を達成し、K=1 の標準的推論手法（2.9 のブロック効率）を大きく上回った。
• K=1 かつ L=8 の条件下で、標準的推論手法よりもさらに 1.37 倍の高速化を達成した。
• 提案されたドラフト選択アルゴリズムは、1−1/e の近似最適受容確率を達成しており、このような手法における最高の乗法的保証である。
• 実験的結果から、SpecTr の出力分布が大規模モデルの真の分布と一致することが確認され、品質劣化がないことが保証された。
• 理論的利点は、実装のオーバーヘッドや非無視可能な小規模モデルの推論時間があっても、実際の実装においてもほとんど保持された。