スループットが約30倍に増加しました。田元東チームの最新論文は、大規模モデル展開の問題を解決している

大規模言語モデル (LLM) は今年非常に人気がありました。しかし、その驚異的な効果の背後には、巨大なモデルと誇張されたハードウェア リソースがあります。

LLM を実際に導入する場合、通常、高価な KV キャッシュ コストと長いシーケンスの一般化能力の低さという 2 つの課題に直面します。

最近、Tian Yuandong 氏のチームは、上記の 2 つの問題を解決し、推論システムのスループットを 30 倍近く向上させた論文を発表しました。

論文アドレス: https://arxiv.org/pdf/2306.14048.pdf

コードアドレス: https://github.com/FMInference/H2O

この成果は NeurIPS'23 でも発表される予定です。

次に、これら 2 つの問題の詳細と、論文で提示されている解決策を見てみましょう。

1 つ目はキャッシュです。KV キャッシュは、再計算を避けるために、生成プロセス中に中間アテンション キーと値を保存するために使用されます。

通常、モデルパラメータに加えて、大量の一時情報（KV キャッシュ）が GPU メモリに保存されます。この部分のメモリ使用量は、シーケンスの長さとバッチ サイズに比例します。

たとえば、入力バッチ サイズが 128、シーケンス長が 1024 の 300 億パラメータ モデルでは、180 GB の KV キャッシュが必要です。

2 番目に、ハードウェアの制限により、LLM は固定のシーケンス長で事前トレーニングされます (たとえば、Llama-2 は固定長の 4K シーケンスを使用します)。

しかし、これにより推論中の注意ウィンドウにも制限が課せられ、より長い入力シーケンスに直面したときにモデルが効果を発揮しなくなり、より幅広いアプリケーションへの適用が妨げられます。

これに対応して、この論文では、メモリ使用量を大幅に削減し、長い入力シーケンスを持つタスクで優れたパフォーマンスを発揮する KV キャッシュを実装する新しい方法を提案しています。

この方法は、注目度スコアを計算する際に少数のトークンがほとんどの値に貢献するという事実に基づいており、これをヘビーヒッター (H2) と呼びます。

包括的な調査により、H2 の出現は自然であり、テキスト内のフレーズの頻繁な共起と密接に関連していることがわかりました。一方、H2 を削除すると、パフォーマンスが大幅に低下します。

これに基づいて、著者らは、最近のトークンと H2 トークンのバランスを動的に維持する KV キャッシュ削除戦略である Heavy Hitter Oracle (H2O) を提案しました。

さらに、著者らは KV キャッシュの削除を動的サブモジュール問題として定式化し、提案された削除アルゴリズムに理論的な保証を提供します。

最後に、著者らは OPT、LLaMA、GPT-NeoX を使用して、さまざまなタスクにおけるアルゴリズムの精度を検証しました。

その中で、OPT-6.7BとOPT-30Bに実装されたH2Oは、DeepSpeed Zero-Inference、Hugging Face Accelerate、FlexGenの3つの推論システムのスループットをそれぞれ29倍、29倍、3倍向上させ、同じバッチサイズでは、H2Oはレイテンシを最大1.9倍削減できます。

論文の詳細

上の図は、LLM 生成で展開されるさまざまな KV キャッシュ戦略の象徴的な図です。左下は H2O フレームワークの概要、右下はさまざまな戦略における精度とメモリ消費の比較です。

事前トレーニング済みの LLM に以前の方法を適用すると、ミス率が高くなり、精度が低下することがわかります。

KV キャッシュ問題を解決するには、3 つの技術的な課題があります。

まず、KV キャッシュのサイズを制限できるかどうかは不明です。原理的には、各デコード手順で、以前のすべてのアテンション キーと値にアクセスする必要がある可能性があります。

第二に、生成精度を維持する最適な排除戦略を決定することは、組み合わせ問題です。

最後に、最適なポリシーをブルートフォースで決定することが可能であっても、実際のアプリケーションに展開することは現実的ではありません。

幸いなことに、著者らは研究を通じていくつかの興味深い結果を発見しました。

キャッシュ サイズが小さい場合のスパース性: 高密度トレーニング中でも、LLM のアテンション マトリックスは推論時に 95% を超えるスパース性を持ちます (図 a)。これは、さまざまな事前トレーニング済み LLM で機能します。

したがって、各生成ステップで、同じ出力トークンをデコードするには KV キャッシュの 5% のみで十分であり、精度を損なうことなく KV キャッシュ サイズを最大 20 倍削減できることがわかります。

ヘビーヒッター（H2）：注目ブロック内のすべてのトークンの累積注目スコアは、べき乗分布に従います（図b）。これは、生成プロセスにおいて極めて重要であり、重量級のトークンである影響力のあるトークンの小さなグループが存在することを示唆しています (H2)。これにより、組み合わせ検索の問題から抜け出し、精度を維持する排除戦略を決定できます。

低コスト戦略による貪欲アルゴリズム: 各デコードステップでローカル統計に基づいて H2 を保持する (前のトークンの注目スコアを合計するだけ) ことは、将来のトークンの注目を考慮するのと同じくらい効果的です (図 d)。

上記に基づいて、著者らはサイズ制約のある KV キャッシュでの LLM 生成プロセスを定義し、上記の特性を利用し、シンプルで低コストのエビクション戦略を使用する Heavy-Hitter Oracle (H2O) を提案しました。

方法と分析

LLM 生成プロセスは、次の 2 つの異なるフェーズで構成されます。

ヒントフェーズ: 入力シーケンスを使用して、LLM トレーニング中に採用されたフォワードパスと同様に、KV キャッシュ (キーと値の埋め込みで構成) を生成します。

トークン生成フェーズ: KV キャッシュを活用して更新し、新しいトークンを段階的に生成します。各生成ステップは、以前に生成されたトークンに依存します。

この論文では、トークン生成段階での KV キャッシュの注意効率を改善し、それによって LLM 推論を高速化することに焦点を当てています。

著者らは、注目クエリ行列 Q とキー行列 K を含む有限の KV キャッシュ サイズを持つ生成プロセスを定義します。

立ち退き戦略:

そして、エビクション戦略を使用した生成プロセス:

次に、精度を損なうことなく KV キャッシュ サイズを削減する可能性について説明します。

上の図では、(a) は事前トレーニング済み OPT モデルにおける注目度のスパース性を表します。(b) は対応する単語 (赤い散布点) とデータ内の単語の共起回数 (灰色の曲線) に対する累積注目度スコアの分布を表します。また、x 軸は語彙内の単語インデックスを表します。(c) は完全な KV キャッシュを備えたベースライン モデルとこの論文のモデル (H2O) のパフォーマンス比較を表します。(d) は完全な KV キャッシュを備えたベースライン モデル、ローカル統計を備えた H2O、グローバル統計を備えた H2O、最新の KV (ローカル) のみを備えたモデルの比較を表します。

クエリ行列 Q とキー行列 K から計算された正規化された注目スコア Softmax 行列が与えられた場合、しきい値は各行の最大値の 1% に設定され、対応するスパース性が計算されます。

次に、Wiki-Text-103 の検証セットで事前トレーニング済みの OPT モデルを使用してゼロショット推論を実行し、注意ブロック内のレイヤーごとのスパース性をプロットし、正規化された注意スコア マトリックスを視覚化します。

結果は下の図の通りです。LLM は密にトレーニングされていますが、結果として得られる注目スコア マトリックスは非常にスパースであり、ほぼすべてのレイヤーでスパース性が 95% を超えています。

アテンション ブロックのスパース性は、次のトークンを生成するときに、以前のすべてのキーと値の埋め込みにアクセスする必要がないことを示し、不要な KV 埋め込みを排除して、生成プロセス中の KV キャッシュの必要性を減らすことができます。

ただし、重要な KV が削除されると、LLM によって生成される順次依存関係により LLM のパフォーマンスが低下する可能性があるため、削除戦略には注意が必要です。

著者は、注目ブロック内のすべてのトークンの累積注目スコアが、下の図に示すようにべき乗分布に従うことを発見しました。

これは、生成プロセスにおいて重要な少数のトークン、つまり前述の Heavy-Hitters (H2) が存在することを示しています。

さらに、各単語の累積注目度スコア (赤い点) は、データ内での共起 (灰色の曲線) と高い相関関係にあります。

上記の現象に基づいて、著者は貪欲な排除戦略を設計しました。

生成中に、トークンの数が割り当てられた KV キャッシュ バジェットを超えると、累積されたアテンション スコア統計とキャッシュ内のローカル トークンに基づいて、重量級のトークンが保持されます。

一般的に言えば、最良の結果を得るには、生成プロセス全体からの統計データを使用する必要がありますが、将来生成されるトークンにアクセスすることは不可能であるため、実際の展開では明らかに実現可能ではありません。

そこで、著者は下図に示す実験を行い、各デコードステップでローカル統計を使用して計算されたローカル H2 は、将来のトークンを考慮した場合 (赤線と青線) とほぼ同じであることを発見しました。

次に著者らは、この動的注意スコア計算（空間的制約付き）を新しい動的サブモジュラー型問題として定義します。

上記のフォームを使用して、KV キャッシュ削除ポリシーを定義します。

上の図は、削除アルゴリズムと説明例を示しています。ここでは、KV キャッシュの予算サイズが 3 であると仮定します。4 番目のデコード手順を完了すると、3 番目のトークンに関連付けられた KV 埋め込みは、累積されたアテンション スコアに基づいて削除されます。削除された KV 埋め込みは、後続のデコード手順ではアクセスできなくなります。

さらに、著者は実際の実装の詳細についても言及しています。たとえば、I/O 効率を確保するために、保存されている KV を削除するときにメモリをスワップするのではなく、新しく追加された KV を直接入力します。

実験結果

この論文の実験では、OPT、LLaMA、GPT-NeoX-20B を含む 3 つの代表的な LLM モデル シリーズが選択されました。

8 つの評価タスクが選択されました：COPA、MathQA、OpenBookQA、PiQA、RTE、Winogrande、XSUM、CNN/Daily Mail。

実験的なハードウェアでは、NVIDIA A100 80GB GPU を使用します。

H2O が採用しているキャッシュ戦略を考慮すると、フル KV キャッシュ (Full) に加えて、ローカル キャッシュ戦略 (Local) もベースライン メソッドとして使用されます。

上記の図と表から、さまざまな KV キャッシュ バジェットの下で、この記事で提案された方法 (H2O) は、さまざまな条件下でのテストにおいてローカル戦略よりも優れていることがわかります。

同時に、KV キャッシュ バジェットが 20% 未満である H2O は、完全な KV 埋め込みモデル (Full) と同等のパフォーマンスを実現し、より困難な長いシーケンス生成タスク、XSUM、CNN/Daily Mail でも優れたパフォーマンスを発揮します。