KV キャッシュを使用して LLM 推論をスケーリングする方法

"ChatGPT または Claude を使用する場合、最初のトークンが 2 番目のトークンよりも表示されるまでに時間がかかるのはなぜですか?それは KV キャッシュのおかげです。テキスト生成サービスの提供遅延を改善することを目的とした小さな最適化トリックの数を考えると、これは非常に印象的です。

LLM は、次のトークンを繰り返し予測し、それを以前に生成されたトークンと元のプロンプトに追加することによってテキストを生成します。通常、因果 LLM はシーケンス内の次の単語を予測するようにトレーニングされます。これは、各入力トークンが Transformer 内の結果として得られる非表示状態にマップされ、さらに次のトークンの予測ベクトルにマップされることを意味します。予測ベクトルには、辞書内のトークンと同じ数の予測があり、最も高い確率で予測を貪欲に見つけることで、次のトークンを予測できます。これは、特定のトークンをデコードするときに、対応する隠れ状態を計算するだけで済み、その他は破棄するだけで済むことを意味します。

この必須の隠れ状態は、入力シーケンスの最後のトークンに対応します。この隠れ状態を計算するには、セルフ アテンション層のそれぞれで、入力シーケンス全体のすべてのキーと値が必要ですが、このシーケンスの最後のトークンのクエリだけが必要です。クエリとすべてのキーの間の内積のソフトマックスを取得することで、最後のトークンのアテンションを生成します。

最後のトークンのアテンション = Softmax(クエリ x すべてのキー)

そして、アテンションによって与えられるすべての値の加重平均を取ることで、結果として得られる隠れた状態を取得します。

非表示の状態 = すべての値 x アテンション

導き出せる結論はいくつかあります。まず、最後のトークンに対応するアテンション以外のアテンションを計算する必要がありません。これは、通常のアテンション メカニズムを使用した場合でも、デコード プロセス中の各反復の時間計算量は入力トークンの数 (~O(N)) に線形であることを意味します。第 2 に、キーと値はすべての反復で各トークンで同じままです。これは、反復ごとにそれらを再計算する必要がないことを意味します。そこから KV キャッシュのアイデアが生まれました。

デコードプロセスは、初期化フェーズと生成フェーズの 2 つのフェーズに分割できます。初期化フェーズでは、入力プロンプト内のすべてのトークンに対応するすべてのキーと値を作成する必要があります。これには、次のフェーズとほぼ同じ時間がかかります。その後、すべてのアテンション レイヤーのすべてのキーと値を KV キャッシュに保存できます。生成フェーズでは、入力シーケンスの最後のトークンに対応するキー、クエリ、および値を生成するだけで済みます。次に、保存されているキーと値をキャッシュから取得して、必要な隠れ状態を計算できます。最後に、最後のトークンに対応する最新の計算されたキーと値でキャッシュを更新します。

この KV キャッシュ プロセスにより、テキスト生成に伴う待ち時間が大幅に短縮されます。"

日本語版は Ai 支援を使用しているため、小さな間違いが存在する可能性があることをご了承ください。

著者

Ai Base Network (ABN), ABN ASIAは、アカデミアに深く関わり、アメリカ、オランダ、ハンガリー、日本、韓国、シンガポール、ベトナムでの仕事経験を持つ人々によって設立されました。ABN ASIAは、学問とテクノロジーが機会と出会う場所です。最先端のソリューションと優れたソフトウェア開発サービスにより、ビジネスがレベルアップし、グローバルシーンに挑戦できるよう支援しています。 私たちの取り組み: より速く。 より良い。 より信頼性が高くなります。 ほとんどの場合、価格も安くなります。

いつでも、ITサービス、デジタルコンサルティング、既製のソフトウェアソリューション、または提案依頼書（RFP）をお探しの際は、お気軽にお問い合わせください。お問い合わせ先は[email protected]です。お客様のテクノロジーに関するニーズにお応えします。

© ABN ASIA