TransformerのKVキャッシュを用いたDecode処理を理解する - Qiita

こんにちは、DeNAでデータサイエンティストをやっているまつけんです。 今回は、TransformerのKVキャッシュを用いたDecode処理について解説していきたいと思います。 Transformerの推論処理には以下のように２段階あります。 Prefill: 入力トークン全体(プロンプト)を並列処理する。またその際、K,V を作成してキャッシュとして保存 Decode: プロンプトの入力から、１トークンずつ処理して文章を生成。過去作成したK,Vをキャッシュから読み出して再利用し、新しい1行(1トークン)を追加 1段階目のPrefillは初回処理として以前の記事↓を参照してもらえればと思います。 この処理を行い出来上がったKとVを保存しておき、これをKVキャッシュと呼びます。 この計算、 $X$を入力として次のレイヤーへの出力$Y$を作る処理になります。 図で表すと以下のようになります。

[misshiki]

TransformerのDecodeは、過去のK/VをKVキャッシュから再利用し、新規1トークン分だけ計算。系列全体の再計算を避けて効率化できる一方、長文ではキャッシュ容量が増大する。