[論文レビュー] SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference

Luciano Del Corro, Allie Del Giorno+4 more|arXiv (Cornell University)|Jul 5, 2023

Topic Modeling被引用数 7

ひとこと要約

SkipDecode は、シーケンスの各位置で単調に減少する終了ポイントを用いる、自己回帰的 LLM 推論における新しいトークンレベルの早期終了手法であり、バッチ処理と KV キャッシュを効率的に行える。低い層をスキップし、上位層に計算リソースを集中させることで、性能の著しい低下を伴わず 2x–5x の高速化を達成する。実用的なデプロイメントパイプラインと互換性がある。

ABSTRACT

Autoregressive large language models (LLMs) have made remarkable progress in various natural language generation tasks. However, they incur high computation cost and latency resulting from the autoregressive token-by-token generation. To address this issue, several approaches have been proposed to reduce computational cost using early-exit strategies. These strategies enable faster text generation using reduced computation without applying the full computation graph to each token. While existing token-level early exit methods show promising results for online inference, they cannot be readily applied for batch inferencing and Key-Value caching. This is because they have to wait until the last token in a batch exits before they can stop computing. This severely limits the practical application of such techniques. In this paper, we propose a simple and effective token-level early exit method, SkipDecode, designed to work seamlessly with batch inferencing and KV caching. It overcomes prior constraints by setting up a singular exit point for every token in a batch at each sequence position. It also guarantees a monotonic decrease in exit points, thereby eliminating the need to recompute KV Caches for preceding tokens. Rather than terminating computation prematurely as in prior works, our approach bypasses lower to middle layers, devoting most of the computational resources to upper layers, allowing later tokens to benefit from the compute expenditure by earlier tokens. Our experimental results show that SkipDecode can obtain 2x to 5x inference speedups with negligible regression across a variety of tasks. This is achieved using OPT models of 1.3 billion and 6.7 billion parameters, all the while being directly compatible with batching and KV caching optimization techniques.

研究の動機と目的

• 自己回帰的 LLM 推論における、従来のトークンレベルの早期終了戦略がバッチ処理と KV キャッシュと互換性がないという問題を解決すること。
• バッチ内の最後のトークンが終了するまで待たなければならないというボトルネックを克服し、高速化の可能性を制限すること。
• バッチ内のすべてのトークンに対して、シーケンスの各位置で統一された終了ポイントを設けることで、計算リソースの予測可能で制御可能な予算を実現すること。
• シーケンス内の後続トークンがより予測可能であるという観察を活用し、下流のトークンに対して計算量を削減すること。
• 戦略的な計算リソースの割り当てにより上位のトランスフォーマー層に集中させることで、生成品質を維持しつつ、推論遅延を著しく短縮すること。

提案手法

• シーケンスの各位置で単調に減少する終了ポリシーを導入し、初期のトークンはより多くの層を使用し、後続のトークンはより少ない層を使用する。これは、損失で測定される不確実性の低下に基づく。
• バッチ内のすべてのトークンに対して、シーケンスの各位置で1つの共通の終了ポイントを定義し、後続のトークンが終了しても、以前のトークンの KV キャッシュを再計算する必要がなくなる。
• 後続のトークンに対して、下位および中位のトランスフォーマー層をスキップし、計算リソースを上位層に再配分することで、初期のトークンの計算結果を活用する。
• シーケンス全体にわたり線形減衰関数を用いて終了位置を定義し、単調性を保ちつつ、キャッシュされたキー・バリュー状態の再計算を回避する。
• 標準的なバッチ処理と KV キャッシュ最適化とシームレスに統合され、実世界のシステムにおける実用的デプロイメントを可能にする。
• 各シーケンス位置における最大終了層を設定することで、計算リソースの予算を制御し、予測不能な全ネットワーク計算を回避する。

(a) OPT-350m on OpenWebText . Average loss per token position shows a strong monotonically-decreasing trend for general text.

実験結果

リサーチクエスチョン

• RQ1 自己回帰的 LLM 推論において、バッチ処理と KV キャッシュと互換性を持つトークンレベルの早期終了戦略は実現可能か？
• RQ2 シーケンスの各位置で単調に減少する終了ポリシーは、計算コストを削減しつつ生成品質を維持できるか？
• RQ3 後続のトークンが先に終了した場合でも、KV キャッシュの再計算を必要とせずに早期終了を適用できるか？
• RQ4 制御された計算リソース予算下で、SkipDecode の速度向上と性能のトレードオフは、既存手法と比較してどのように異なるか？
• RQ5 バッチ処理と KV キャッシュに依存する実用的なデプロイメントパイプラインを、推論効率を損なわずサポートできるか？

主な発見

• SkipDecode は、OPT-1.3B および OPT-6.7B モデルにおいて、複数の生成ベンチマークで、完全な推論と比較して 2x から 5x の高速化を達成した。
• この手法は、完全なモデルとほぼ同一の生成品質を維持しており、高速化率が高くても性能の著しい低下は見られなかった。
• 単調に減少する終了ポイントを強制することで、SkipDecode は、初期のトークンの KV キャッシュを再計算する必要がなくなり、効率的なバッチ処理が可能になった。
• このアプローチは、バッチ処理と KV キャッシュと完全に互換性があり、従来の早期終了手法の主要な制限を克服した。
• 線形減衰関数による終了位置のポリシーは効果的で頑健であり、初期の実験ではパワーロウ比などの代替関数を上回った。
• オラクル実験により、観察された損失の傾向（トークン位置が進むにつれて減少）が、初期の深層計算を採用する設計選択を裏付けた。

(b) OPT-350m (finetuned) on Reddit-TLDR . Average loss per token position. Decreasing trend but with a different function.