著者 | ピエール・リアンハート 編纂者：岳陽 🚢🚢🚢AIテクノロジーソフトウェアと技術交流グループへのご参加をお待ちしております！最新のトレンドを把握し、一緒に技術的な課題を探求しましょう！ 前回の記事では、Transformer デコーダーのテキスト生成アルゴリズムの概要を説明し、テキスト生成の 2 つのフェーズ、つまりプロンプトを処理するシングルステップの開始フェーズと、一度に 1 つのトークンずつテキストが生成されるマルチステップの生成フェーズに重点を置いて説明しました。この記事では、シーケンス全体 (プロンプト トークンと生成されたテキスト トークンを含む) を各生成ステップの入力として使用する場合に必要な冗長な計算を示します。言い換えると、シーケンス全体を各トークン生成の入力として使用すると不要な計算が発生する可能性があり、この記事では KV キャッシュと呼ばれる手法を使用してこれらの冗長な計算を回避する方法を探ります。簡単に言うと、この手法では、通常は再計算が必要な部分を保存して再利用します。最後に、KV キャッシュが生成フェーズをどのように変更し、開始フェーズとどのように区別するかを理解します。 01 Transformerの注目層について簡単に解説まず、Transformer モデルの最もオリジナルなバージョンのマルチヘッド アテンション (MHA) レイヤーのいくつかの側面を見てみましょう (図 1)。 図 1 — 入力シーケンスの長さが 3 の Transformer デコーダー層 (上) とデュアルヘッド (自己) 注意層 (下) の詳細図。 簡単にするために、長さ t の単一の入力シーケンスのみを処理すると仮定します (つまり、バッチ サイズは 1)。
次に、シングルアテンションヘッドについて説明します。
しかし、自己回帰復号法の文脈では、与えられたクエリベクトルの出力表現を構築するために、すべての可能な値ベクトルを使用することはできません。実際、特定のトークンに関連付けられたクエリベクトルの出力を計算する場合、シーケンス内で後から出現するトークンの値ベクトルを使用することはできません。この制限は、マスキングと呼ばれる手法によって実現されます。マスキングは、基本的に禁止された値ベクトル（つまり、禁止されたトークン）の注目度スコアを0に設定します。 02 マスキング技術を使用すると、生成段階で冗長な計算が発生します。さて、問題の核心について議論する必要があります。マスキング技術の使用により、現在のトークンの出力表現を生成する際に、以前に生成されたトークンの情報のみが使用され、後続のトークンの情報は使用されません。以前のトークンは各反復処理で同じであるため、特定のトークンの出力表現も後続のすべての反復処理で同じになり、冗長な計算が発生します。 前回の記事で紹介したトークンシーケンス（各単語がトークンで構成されているシーケンス）を例に挙げてみましょう。入力シーケンス「What color is the sky? The sky」から「is」を生成したとします。前回の反復処理では、「sky」が入力シーケンスの最後のトークンであったため、このトークンに関連付けられた出力表現は、シーケンス内のすべてのトークンの表現、つまり「What」、「color」、「is」、「the」、「sky」、「?」」、「The」、「sky」の値ベクトルを用いて生成されました。 次の反復処理の入力シーケンスは「空は何色ですか？ 空は」となりますが、マスク処理により、「空」の視点からは、入力シーケンスは依然として「空は何色ですか？ 空」であるように見えます。したがって、「空」に対して生成される出力表現は、前の反復処理の表現と全く同じになります。 ここで、図1の図を用いて例を示します（図2）。この例では、初期化ステップで長さ1の入力シーケンスを処理すると仮定します。著者らは、計算において冗長性を生み出す要素を色分けして強調表示しています。薄い赤と薄い紫は、それぞれ冗長的に計算されるキーベクトルと値ベクトルを表しています。 図2 — 生成フェーズにおける注意層での冗長計算 前の例に戻ると、自己回帰デコードステップの新たな反復では、「空は何色ですか？ 空は」が入力シーケンスとして使用されます。前のステップで計算されていないのは、入力シーケンスの最後のトークン「is」の表現だけです。 もっと具体的に、これを達成するには何をする必要がありますか?
キーと値のベクトルについては、「is」を除き、以前の反復処理で全てのトークンに対して既に計算済みです。したがって、以前の反復処理で得られたキーと値のベクトルを保存（つまりキャッシュ）して再利用することができます。この最適化は単にKVキャッシュと呼ばれます。「is」の出力表現のコンパイルは非常に簡単になります。
この最適化手法では、キーバリューベクトルを使用できる限り、以前のトークンは不要になります。KVキャッシュを使用する場合、モデルへの実際の入力は、最後に生成されたトークン（シーケンス全体ではありません）とKVキャッシュです。下の図3は、生成フェーズでアテンション層を実行するこの新しい手法を示しています。 図3 — KVキャッシュを有効にする生成手順 前の記事で述べた 2 つの段階に戻ります。
アテンションフェーズでは、アテンション層はデコードステップのように一度に1つのトークンを処理するのではなく、すべてのプロンプトトークンを一度に処理します。文献[1]では、最初の設定はバッチアテンション（誤って並列アテンションと呼ばれることもあります）と呼ばれ、2番目の設定はインクリメンタルアテンションと呼ばれています。 キーバリューキャッシュを使用する場合、スタートアップフェーズでは、キーバリューキャッシュ内のすべての入力トークンのキーと値のベクトルを計算し、（事前）設定します。そのため、このフェーズはしばしば事前設定フェーズと呼ばれます。実際には、「事前設定フェーズ」と「スタートアップフェーズ」という用語は同じ意味で使用されますが、ここでは前者を使用することにします。 03 HuggingFace Transformersを使ったキーバリューキャッシュの実装例KVキャッシュは実際にどれほど効果的なのでしょうか？KVキャッシュを有効または無効にすることはできるのでしょうか？HuggingFace Transformers[3]ライブラリを例に挙げてみましょう。テキスト生成専用のすべてのモデルクラス（XXXForCausalLMクラスなど）は、generateというメソッドを実装しており、これは生成プロセス全体の初期エントリポイントとして使用されます。このメソッドは多数の設定パラメータ[4]を受け入れ、主にトークンの検索戦略を制御するために使用されます。KVキャッシュの有効化/無効化は、ブールパラメータuse_cache（デフォルトはTrue）によって制御されます。 さらに一歩進んで、モデルのforwardメソッドを調べると（例えば、LlamaForCausalLM.forward[5]のドキュメントを参照）、use_cacheブールパラメータが簡単に見つかります。KVキャッシュを有効にすると、2つの入力、つまり最後に生成されたトークンとKVキャッシュが、それぞれパラメータinput_idsとpast_key_valuesを介して渡されます。新しいKV値（つまり、現在の反復で計算された新しいキーベクトルと値ベクトル）は、forwardメソッドの出力の一部として返され、次の反復で使用されます。 では、返されるキーと値のペアはどのようなものになるでしょうか？テンソル計算を行ってみましょう。キーと値のキャッシュを有効にすると、`forward` メソッドはテンソルペアのリスト（キーベクトル用と値ベクトル用）を返します。テンソルペアの数は、モデル内のデコーダーブロック（デコーダーレイヤーとも呼ばれ、`n_layers` と表記されます）の数と等しくなります。バッチ内の各シーケンスの各トークンについて、各アテンションヘッドは `d_head` 次元のキー/値ベクトルを持ちます。したがって、各キー/値テンソルの形状は (batch_size, seq_length, n_heads, d_head) となります。 具体的には、MetaのLlama2–7B[6]を例にとると、n_layers=32、n_heads=32、d_head=128となります。KVキャッシュのサイズについては次の記事で詳しく説明しますが、現時点では実現可能なサイズについて、ある程度の直感的な理解が得られています。 04 KV キャッシュを使用すると計算量をどれくらい節約できますか?入力シーケンスのバッチbがあるとします。各シーケンスはN個の生成トークンとt個の入力トークン（合計長N+t）で構成されています。これらのシーケンスの最初のt+N-1個のトークンについては、キー値（KV）の計算は冗長です。つまり、N番目の生成ステップでは、各シーケンスについてt+N-1回のKV計算を節約できます。再計算を行わない場合、最初のN番目の生成ステップでは、各シーケンスについて合計N.t+N.(N-1)/2回のKV計算を節約できます。 ステップNで再コンパイルを行わない場合、どれだけのFLOPsを節約できるでしょうか？特定のトークンのキーベクトルまたは値ベクトルを計算するには、サイズd_modelの埋め込みベクトルと形状(d_model, d_head)の重み行列を掛け合わせるだけで済みます。これをさらに詳しく見ていきましょう。 各生成ステップにおいて、不要なキーまたは値ベクトルの計算はいくつ実行されるでしょうか？各デコーダー層では、各トークンと各アテンションヘッドに対して2回の計算（キーベクトルと値ベクトルのそれぞれ1回ずつ）が実行されます。つまり、各トークンには2.b.n_layers.h回の計算が必要です。つまり、最初のN生成ステップでは、各シーケンスによって合計b.n_layers.hN(t+N-1)回のキーまたは値ベクトルの計算が節約されます。 1回の行列乗算には何FLOPS必要ですか？形状(n, p)の行列とサイズ(n, m)の別の行列を乗算するには、約2.mnp回の演算が必要です。したがって、この例では、キーベクトルまたは値ベクトルを計算するには、約2.d_model.d_head回の演算が必要です。 全体的に、KV キャッシュを選択すると、最初の N 生成ステップでおよそ次の数の FLOP が節約されます。 キーバリューキャッシュを使用することで、最初のt+N-1個のトークンに対するクエリベクトルの計算を回避し、t+N-1個の出力表現に出力重み行列を乗算する必要もなくなります。ただし、上記の式は変更しません。最初のt+N-1個のトークンに対する注目度スコアを計算しないことで、以下のFLOPを節約できます。 実際の数値としては、MetaのLlama2-7B[6]を例にとると、n_layer=32、d_model=128、d_model=4096となる。 最も重要なのは、KVキャッシュによる計算コストの削減は、生成されるトークンの数の2乗に比例することに注目したことです。（訳注：つまり、生成されるトークンの数が2倍になると、KVキャッシュによる計算コストの削減は4倍になります。） ただし、ここまではKVキャッシュのメリットのみを見てきました。デメリットについては次の記事で解説します。KVキャッシュは妥協であり、タダ飯ではないことを覚えておいてください。つまり、メモリとデータ転送量の増加と引き換えに、計算量を削減しているのです。次の記事で説明するように、KVキャッシュの使用コストは相当に高くなる可能性があり、他のトレードオフと同様に、必ずしも費用対効果が高いとは限りません。 05 結論最後に、この記事で学んだことをまとめましょう。アテンション計算ではマスキングを使用しているため、各生成ステップにおいて、過去のトークンのキーと値のベクトルを再計算する必要はなく、最終的に生成されるトークンのキーと値のベクトルを計算するだけで済みます。新しいキーと値のベクトルを計算するたびに、それらをGPUメモリにキャッシュして将来再利用できるため、再計算にかかる計算コストを節約できます。 このキャッシュ戦略を適用すると、生成フェーズにおけるアテンション層への入力が、起動フェーズで必要な入力と比較して変化します。起動フェーズでは、アテンション層は入力シーケンス全体を一度に処理しますが、KVキャッシュを有効にした生成フェーズでは、最後に生成されたトークンとKVキャッシュのみが入力として必要になります。起動フェーズと生成フェーズの間のこの新しい違いは、単なる概念的なものではありません。例えば、各フェーズで特定のGPUカーネルを使用すると、両フェーズで同じGPUカーネルを使用するよりもパフォーマンスが向上する可能性があります[2]。 上で述べたように、キー値キャッシュは簡単に得られるものではなく、一連の新たな問題を引き起こします。これについては、次の記事で詳しく説明します。
次の記事では、KVキャッシュサイズの問題について考察します。その後の記事では、ハードウェアの利用率についてより詳細に解説し、KVキャッシュを省略できる状況についても考察します。 読んでくれてありがとう！ 🚢🚢🚢AIテクノロジーソフトウェアと技術交流グループへのご参加をお待ちしております！最新のトレンドを把握し、一緒に技術的な課題を探求しましょう！ 終わり
|
