LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

要約

トランスフォーマーベースの大規模言語モデルの推論は、2 つの連続したステージで構成されます。1) プロンプトの KV キャッシュを計算して最初のトークンを生成するプレフィル ステージと、2) 後続のトークンを生成するデコード ステージです。
長いプロンプトの場合、事前入力段階ですべてのトークンに対して KV キャッシュを計算する必要があるため、最初のトークンの生成に必要な時間が大幅に長くなる可能性があります。
その結果、プレフィル段階が生成プロセスのボトルネックになる可能性があります。
最初のトークンを生成するためにすべてのプロンプト トークンが必須であるかどうかについては、未解決の疑問が残っています。
これに答えるために、プレフィルとデコードの両方の段階で次のトークンの予測に重要なトークンの KV を選択的に計算する新しいメソッド LazyLLM を導入します。
プロンプトを一度にプルーニングする静的プルーニング アプローチとは対照的に、LazyLLM では、たとえ前のステップでプルーニングされたとしても、言語モデルがさまざまな生成ステップでコンテキストからトークンのさまざまなサブセットを動的に選択できます。
さまざまなタスクにわたる標準データセットに対する広範な実験により、LazyLLM が既存の言語モデルとシームレスに統合でき、微調整することなく生成を大幅に加速できる汎用メソッドであることが実証されました。
たとえば、複数ドキュメントの質問応答タスクでは、LazyLLM は精度を維持しながら、LLama 2 7B モデルの事前入力ステージを 2.34 倍高速化します。

要約(オリジナル)

The inference of transformer-based large language models consists of two sequential stages: 1) a prefilling stage to compute the KV cache of prompts and generate the first token, and 2) a decoding stage to generate subsequent tokens. For long prompts, the KV cache must be computed for all tokens during the prefilling stage, which can significantly increase the time needed to generate the first token. Consequently, the prefilling stage may become a bottleneck in the generation process. An open question remains whether all prompt tokens are essential for generating the first token. To answer this, we introduce a novel method, LazyLLM, that selectively computes the KV for tokens important for the next token prediction in both the prefilling and decoding stages. Contrary to static pruning approaches that prune the prompt at once, LazyLLM allows language models to dynamically select different subsets of tokens from the context in different generation steps, even though they might be pruned in previous steps. Extensive experiments on standard datasets across various tasks demonstrate that LazyLLM is a generic method that can be seamlessly integrated with existing language models to significantly accelerate the generation without fine-tuning. For instance, in the multi-document question-answering task, LazyLLM accelerates the prefilling stage of the LLama 2 7B model by 2.34x while maintaining accuracy.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google