推論 — 知識ノート

学習済みモデルで実際にテキストを生成する処理。inference。

自己回帰生成#

1 トークンずつ生成し、それを入力に足して次を生成する。

入力 → 次トークンの確率分布 → サンプリング → 入力に追加 → 繰り返し

逐次的なので、系列長に比例した時間がかかる。訓練時は全位置を並列に計算できるが、生成時はできない。

各ステップで過去すべてを再計算するのは無駄。 Attention の $K$ と $V$ を保存して使い回す。

𝒪 (T^{2}) \to 𝒪 (T) (1 トークンあたり)

代わりにメモリを消費する。

KV キャッシュ = 2 \times L \times T \times d \times バッチ

長文脈・大バッチではこれが GPU メモリの大半を占め、同時に処理できるリクエスト数を制限する。

段階	内容	ボトルネック
Prefill	入力プロンプトを一括処理	計算律速
Decode	1 トークンずつ生成	メモリ帯域律速

Decode では、1 トークン生成するのに全パラメータを読み出す必要がある。計算量は小さいのにメモリ帯域が限界になる。

これが推論の高速化で量子化（パラメータを小さくする）や投機的デコーディング（小モデルで下書きして大モデルで検証）が効く理由。

Dan Jurafsky, James H. Martin. Speech and Language Processing, 3rd ed. draft.（全文公開） https://web.stanford.edu/~jurafsky/slp3/
Reiner Pope et al. Efficiently Scaling Transformer Inference. MLSys, 2023. https://arxiv.org/abs/2211.05102
Yaniv Leviathan, Matan Kalman, Yossi Matias. Fast Inference from Transformers via Speculative Decoding. ICML, 2023. https://arxiv.org/abs/2211.17192