SpecTr: Fast Speculative Decoding via Optimal Transport

要約

大規模な言語モデルからの自己回帰サンプリングにより、いくつかの自然言語タスクで最先端の結果が得られました。
ただし、自己回帰サンプリングでは一度に 1 つずつトークンが生成されるため、処理が遅くなり、特定のタスクでは法外に遅くなる場合もあります。
サンプリングを高速化する 1 つの方法は、$\textit{投機的デコード}$ です。小規模なモデルを使用して $\textit{ドラフト}$ (トークンのブロックまたはシーケンス) をサンプリングし、その後、ドラフト内のすべてのトークンを大きな言語でスコア付けします。
並行してモデル化します。
最終出力が大規模モデルの分布に従っていることを保証する統計的手法に基づいて、ドラフト内のトークンのサブセットが受け入れられます (残りは拒否されます)。
この研究では、$\textit{会員費用}$ を使用した最適トランスポート (OT) のレンズを通して、投機的デコードの原理的な理解を提供します。
このフレームワークは、よく知られている $\textit{maximal-coupling}$ 問題の拡張として見ることができます。
この新しい定式化により、投機的復号方法を一般化して、トークン レベルで $k$ 候補のセットを可能にし、最適なメンバーシップ コストの向上につながります。
最適なドラフト選択アルゴリズム (輸送計画) が線形計画法によって計算できることを示します。その最もよく知られている実行時間は $k$ で指数関数的です。
次に、合格確率が $(1-1/e)$-optimal となる有効なドラフト選択アルゴリズムを提案します。
さらに、単一トークンのドメインのサイズにほぼ直線的に時間内に計算できます。
この $newdraftselection$ アルゴリズムを使用して、$\textit{SpecTr}$ と呼ばれる新しい自己回帰サンプリング アルゴリズムを開発します。これにより、デコードされた出力の品質が低下しないようにしながら、デコードの速度が向上します。
最先端の大規模言語モデルの場合、提案されたアプローチは、標準ベンチマークでの投機的デコードと比較して、実時間で 2.13 倍、さらに 1.37 倍の高速化を達成することを実験的に示します。

要約(オリジナル)

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 $\textit{speculative decoding}$: use a small model to sample a $\textit{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 $\textit{membership cost}$. This framework can be viewed as an extension of the well-known $\textit{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 $\textit{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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google