本文にスキップ
AI News HubLIVE
サイト内リライト4 分で読了

プロンプトを入力してから最初の単語が表示されるまでの仕組み

記事の要約

大規模言語モデルの推論プロセスを詳しく解説。自己回帰生成、プリフィルとデコードフェーズ、KVキャッシュの仕組みとメモリ使用量、デコーディング戦略が出力に与える影響について説明する。

ソースHacker News AI著者: shbhmrzd
プロンプトを入力してから最初の単語が表示されるまでの仕組み
誤りを報告

訂正窓口はまだ利用できません。記事情報をコピーして保存できます。

訂正案内
本文へ

本記事は、LLM(大規模言語モデル)の仕組みをソフトウェアエンジニア向けに解説するシリーズの第4回であり、学習済みモデルが実際に推論を行う際の内部動作に焦点を当てる。ユーザーがプロンプトを入力しEnterキーを押してから最初の単語が表示されるまでの間、モデル内部では何が起こっているのか、またなぜ応答が一度にすべて出力されずに1ワードずつストリームされるのかを説明する。

自己回帰生成

トレーニング時にはモデルはシーケンス全体を一度に処理し、各位置で次のトークンを並列に予測する。しかし推論時には、まだ生成されていない応答を参照することはできないため、モデルは1トークンずつ応答を構築する必要がある。具体的なステップは次の通り:

  1. ユーザーがプロンプトを入力 → トークナイザがトークンIDに変換。
  2. プロンプト全体を入力として1回の前向きパスを実行 → モデルはプロンプト全体の直後に来るトークンのロジットを出力。
  3. ロジットから1つのトークンをサンプリング。
  4. そのトークンをシーケンスに追加。
  5. 元のプロンプト+新しいトークンを入力として再び前向きパス。
  6. 次のトークンをサンプリング。
  7. 希望の長さに達するか停止トークンに当たるまで繰り返す。

これが自己回帰生成であり、各ステップが次のステップの入力となるため、逐次的に行わざるを得ない。

例として、「What is gravity?」というプロンプトは[What, is, gravity, ?](4トークン)に変換される。ステップ1ではこれら4トークンを入力、前向きパス後「Gravity」がサンプリングされる。ステップ2では入力が[What, is, gravity, ?, Gravity](5トークン)となり「is」がサンプリングされる。以下同様に「Gravity is the force…」と生成が続く。

プリフィルとデコード

実際の推論は2つのフェーズに分けられる。

プリフィルフェーズ:プロンプト全体を一度に並列処理する。50トークンのプロンプトなら1回の前向きパスで完了。GPUはすべてのトークンをすべてのレイヤーで同時に処理する(因果マスクにより未来のトークンへの注意は防止される)。プリフィルの真の目的は注意機構の計算であり、各トークンが過去のすべてのトークンに注意を払う過程でKeyとValueのベクトルが生成される。これらのベクトルはそのままKVキャッシュに格納され、デコードフェーズで再利用される。短いプロンプトは高速だが、数万トークンに及ぶ長いプロンプトでは注意機構の計算量がシーケンス長の二乗に比例するため、プリフィル自体が顕著な時間を要することがある。

デコードフェーズ:新しいトークンを1つずつ逐次的に生成する。各前向きパスでは1つの新しいトークンのみを処理するため、GPUが行う計算はごく小規模である。計算のボトルネックはメモリ帯域幅に移る——GPUは演算よりもモデルの重みをメモリから読み込むのにほとんどの時間を費やす。これが、デコード性能においてGPUのメモリ帯域幅が生の計算速度よりも重要である理由である。典型的なチャットの場合、100トークンのプロンプトと500トークンの応答では、プリフィル(1回の前向きパス)は一瞬で終わるが、デコード(500回の逐次前向きパス)がユーザーが感じるレイテンシとなる。応答が1ワードずつストリームされるのは、各トークンが実際に1つずつ生成され、準備ができ次第送信されるためである。

KVキャッシュ:ステップ間の作業を節約

デコードでは各ステップでそれまでの全トークンを処理する。例えばステップ100ではトークン1〜99はステップ99から変わっていないにもかかわらず、注意機構がそれらのKeyとValueベクトルを再計算する。KVキャッシュはこの無駄をなくす:各トークンのKとVベクトルは一度計算されると保存され、以降のステップでは再利用される。プリフィルでプロンプト全体のKとVがキャッシュされ、デコードでは新しいトークンのKとVのみを計算し、既存のキャッシュと連結して注意を実行する。

KVキャッシュのメモリ使用量

Llama 3.1-8Bの場合:32レイヤー、8個のKVヘッド(グループ化クエリ注意機構による)、ヘッド次元128、FP16精度(1値あたり2バイト)、最大コンテキスト長128,000トークン。1トークンあたりのKVキャッシュメモリ:1レイヤー、1ヘッドでKeyベクトル128×2=256バイト、Valueも同様で合計512バイト。8ヘッドで512×8=4,096バイト、32レイヤーで4,096×32=131,072バイト≈128KB。4,096トークンのシーケンスでは128KB×4,096=512MB。128,000トークンでは約16GBになる。複数ユーザーを処理する場合はさらに増加し、32ユーザー各4,096トークンで合計16GBとなり、モデル重み(約16GB)と同等のメモリを消費する。グループ化クエリ注意機構(GQA)はKVヘッド数を32から8に減らすことでキャッシュサイズを1/4にしており、メモリ節約の効果は大きい。

KVキャッシュの節約効果

キャッシュなしで1,000トークンを生成する場合、K/V計算は1+2+…+1,000≈500,000回必要。キャッシュを使えば各トークンのK/Vは一度だけ計算されるため、合計1,000回で済む。ただし注意のドット積の数はキャッシュの有無にかかわらず約500,000回(O(n²))であり、キャッシュは主にK/Vベクトル計算の行列乗算を削減する。

デコーディング戦略

モデルが各トークンの確率を出力した後、実際にどのトークンを選択するかを決定する戦略が必要である。代表的なものに、グリーディデコーディング(常に最も確率の高いトークンを選ぶ)、温度サンプリング(確率分布の尖り具合を調整)、top-kサンプリング(確率上位k個のトークンのみからサンプリング)、top-pサンプリング(累積確率がpを超える最小のセットからサンプリング)などがある。これらの戦略は応答の創造性と一貫性のバランスを調整し、チャットでは高めの温度、事実回答では低めの温度またはグリーディが適している。

以上のように、LLMの推論は複雑だが巧妙なプロセスであり、プロンプト入力から最終応答までの各ステップに多くの最適化が施されている。プリフィル、デコード、KVキャッシュ、デコーディング戦略を理解することで、開発者はこれらのモデルをより効果的に活用・最適化できるようになる。

要点と分析を開く

記事インテリジェンス

エンジニア上級

要点

  • LLMの推論は自己回帰的で、トークンは1つずつ生成され、各ステップは前の出力に依存する。
  • プロセスは高速なプリフィルフェーズ(プロンプト全体を並列処理)と低速なデコードフェーズ(トークンを逐次生成)に分かれる。
  • KVキャッシュは過去のトークンのKeyとValueベクトルを保存し、冗長な計算を大幅に削減するが、大きなメモリを消費する。
  • 温度やtop-kサンプリングなどのデコーディング戦略が、応答の創造性と一貫性を調整する。

要点と分析は自動生成され、誤りを含む場合があります。原典をご確認ください。