PaSS: Parallel Speculative Sampling

要約

言語モデルのサイズを数百億のパラメーターに拡張することで、幅広いタスクで優れたパフォーマンスが実現しました。
これらのモデルは生成時に自動回帰的に使用され、生成されたトークンごとに前方パスが必要となるため、パラメーターの完全なセットがメモリから読み取られます。
このメモリ アクセスは生成の主なボトルネックを形成し、モデル サイズが大きくなるにつれて悪化します。
さらに、複数のトークンのフォワード パスを並行して実行すると、1 つのトークンのみの場合とほぼ同じ時間がかかることがよくあります。
これら 2 つの観察は、投機的サンプリングの開発につながります。そこでは、2 つ目の小さなモデルを使用していくつかのトークンをドラフトし、その後、大きなモデルの 1 回の前方パスを使用して検証または拒否されます。
残念ながら、この方法では同じトークナイザーを共有する 2 つのモデルが必要なため、その採用が制限されます。
代替案として、計算コストや 2 番目のモデルを必要とせずに、単一のモデルから複数のトークンをドラフトする方法として並列デコードを使用することを提案します。
私たちのアプローチでは、同時に生成される単語をマークする追加の入力トークンのみが必要です。
$O(d_{emb})$ の追加パラメータのみを必要としながら、有望なパフォーマンス (最大 $30\%$ の高速化) を示しています。

要約(オリジナル)

Scaling the size of language models to tens of billions of parameters has led to impressive performance on a wide range of tasks. At generation, these models are used auto-regressively, requiring a forward pass for each generated token, and thus reading the full set of parameters from memory. This memory access forms the primary bottleneck for generation and it worsens as the model size increases. Moreover, executing a forward pass for multiple tokens in parallel often takes nearly the same time as it does for just one token. These two observations lead to the development of speculative sampling, where a second smaller model is used to draft a few tokens, that are then validated or rejected using a single forward pass of the large model. Unfortunately, this method requires two models that share the same tokenizer and thus limits its adoption. As an alternative, we propose to use parallel decoding as a way to draft multiple tokens from a single model with no computational cost, nor the need for a second model. Our approach only requires an additional input token that marks the words that will be generated simultaneously. We show promising performance (up to $30\%$ speed-up) while requiring only as few as $O(d_{emb})$ additional parameters.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google