こんにちは、DeNAでデータサイエンティストをやっているまつけんです。

今回は、TransformerのKVキャッシュを用いたDecode処理について解説していきたいと思います。

Transformerの推論処理には以下のように２段階あります。

Prefill: 入力トークン全体(プロンプト)を並列処理する。またその際、K,V を作成してキャッシュとして保存
Decode: プロンプトの入力から、１トークンずつ処理して文章を生成。過去作成したK,Vをキャッシュから読み出して再利用し、新しい1行(1トークン)を追加

1段階目のPrefillは初回処理として以前の記事↓を参照してもらえればと思います。

この処理を行い出来上がったKとVを保存しておき、これをKVキャッシュと呼びます。

この計算、 $X$を入力として次のレイヤーへの出力$Y$を作る処理になります。

\begin{aligned} 
Q &= X W_Q \cr
K &= X W_K \cr
V &= X W_V \cr 
\end{aligned}

Y = \mathrm{softmax}\left(\frac{QK^{\mathrm{T}}}{\sqrt{d}} + M\right)\cdot V

図で表すと以下のようになります。¹

ここで、$n$はtokenの数、$d$はtokenの次元を表します。

この計算は、Decoder onlyのTransformerで言うと[Multi-Head Attention]の各ボックスで行われ、 X を入力して Y が出力されます。²

Prefill

さて、以上の処理は初回は上記に書いた通り、入力 $X$が基本複数のtokenになることが圧倒的なので、$X$は$n \times d$の行列です。ここまでは前回の記事で書いた内容です。つぎの Decode からが今回の記事のテーマ、KVキャッシュが効いてくるところです。

Decode

Prefillで入力されたプロンプトを処理し、最初の生成トークンを決定します。その後のDecodeでは、直前に生成したトークンを1つずつモデルへ入力し、次のトークンを自己回帰的（auto-regressive）に生成します。

概念上は、生成したトークンを元の入力系列の末尾に追加して処理を繰り返していると考えられます。ただし実際には、毎回入力系列全体を結合して再計算すると、全く同じ処理を繰り返すと言う不要なコストがかかってしまうので、新しいトークンに対応する (Q,K,V) だけを計算し、過去の (K,V) はKVキャッシュから再利用します。これがKVキャッシュです。

例えば"[BOS] two kids are playing in a swimming pool with a green colored crocodile." を入力して次が"Then"だったとしましょう。

n個のtokenが既に計算済みなので、n+1個目のトークンから$q_{n+1}, k_{n+1}, v_{n+1}$の３つのベクトルを作ります。

次に、QKVの計算をしていきますが、Kはnまでのデータはキャッシュを利用し、$k_{n+1}$を下に追加します。同様にVもnまでのデータはキャッシュを利用し、$v_{n+1}$を下に追加します。

$q_{n+1}$と、キャッシュされたKに$k_{n+1}$をつなげた新しい$K$を掛け合わせ、$s_{n+1}$を得ます。

$s_{n+1}$にSoftmaxをかけて、さらにキャッシュされたVに$v_{n+1}$をつなげた新しい$V$を掛け合わせ、$y_{n+1}$を得ます。

これで、このLayerで必要とされていた$y_{n+1}$を得ることができました！

上記では逐次計算しましたが、行列での概念的には以下のように、計算済みの部分、未来参照なので計算不要なCausal Mask部分、今回計算した部分、に分けて考えることもできます。

まとめて表現すると以下の通りです。

$y_{n+1}$はこれだけでいいのか、$y_{0}$〜$y_{n}$は必要ないのか、と思いますが、KとVがキャッシュされていれば$x_{n+1}$だけあれば計算ができたことを思い出すと、これで十分であることがわかります。

以上から、KとVがキャッシュされていれば、以降は次のトークンのベクトルさえ得られればnext tokenの出力ができることがわかりました。非常に効率的ですね！

このKVキャッシュは、Transfomer Layerごとに存在します。上記の図ではN層のLayerがありますが、このN層ごとにKVキャッシュがあるため、$(K_{\mathrm{cache}}^{(1)}, V_{\mathrm{cache}}^{(1)}), \ldots, (K_{\mathrm{cache}}^{(N)}, V_{\mathrm{cache}}^{(N)})$
のように、N個のキャッシュが存在します。

また、この計算には、問題もあります。token数（コンテキストウィンドウ長）が多くなるとK, Vのキャッシュの量が増大していってしまうことです。これに対する対策も考えられていますが、それはまた別の機会で解説したいと思います！

本記事では図を簡単にするためscale係数$\sqrt{d}$で割ることを省略しています ↩
以降ではKVキャッシュの仕組みに集中するため、Multi-Head Attentionを1つのAttention headとして単純化して表現します。 ↩