キャッシュの仕組み

誰かと会話している状況を想像してみてください。相手が以前に言ったことを覚えておらず、あなたが毎回応答するたびに最初からやり直さなければならないとします。これは遅くて非効率ですよね？

この類推は、transformerモデルにも拡張できます。自己回帰モデルによる生成は、トークンごとに1つずつ予測を行うため、遅くなる場合があります。各新しい予測は、それまでのすべてのコンテキストに依存しています。

1000番目のトークンを予測するには、モデルは直前の999トークンの情報が必要です。この情報は、トークン表現間の行列積として表現されます。

1001番目のトークンを予測するには、1000番目のトークンの情報に加えて、それ以前の999個のトークンからの同じ情報が必要です。これは、モデルが各トークンに対して繰り返し計算しなければならない大量の行列積です！

キーバリュー（KV）キャッシュは、以前に処理されたトークンのアテンション層から導出されたkvペアを保存することで、この非効率性を解消します。保存されたkvペアはキャッシュから取得され、後続のトークンに再利用されるため、再計算の必要がなくなります。

は推論でのみ使用すべきです。学習中に有効にすると、予期しないエラーが発生する可能性があります。

がどのように、そしてなぜ機能するのかをより深く理解するために、アテンション行列の構造を詳しく見ていきましょう。

アテンション行列

スケーリングドットプロダクトアテンション は、バッチサイズ 、アテンションヘッド数 b, 、これまでのシーケンス長 h, 、およびアテンションヘッドごとの次元数 T, に対して、以下のように計算されます。 d_head. $\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$

クエリ(Q)、キー(K)、バリュー(V)行列は、形状(b, h, T, d_head)の入力埋め込みからの射影です。

因果的アテンションでは、マスクによってモデルが将来のトークンにアテンションを向けることを防ぎます。トークンが処理されると、その表現は将来のトークンに対して決して変化しないため、 $K_{\text{past}}$ と $V_{\text{past}}$ をキャッシュして再利用し、最後のトークンの表現を計算できます。 $\text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{cached}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{cached}}, v_{t}])$

推論時では、次のトークン $x_t$ を予測する表現を計算するために、最後のトークンのクエリのみが必要です。各ステップで、新しいキーとベクトルは キャッシュに 保存され、 過去のキーとベクトルに 追加されます。 $K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$

Attentionはモデルの各層で独立して計算され、キャッシュは層ごとに実行されます。

が効率をどのように向上させるかを比較するには、以下の表を参照してください。

キャッシュなし	キャッシュを使用
各ステップで、それ以前のすべての`K`と`V`を再計算します	各ステップで、現在の`K`と`V`のみを計算します
ステップあたりのattentionコストは、シーケンス長に対して二次の比例関係にあります	ステップあたりのattentionコストはシーケンス長に対して線形です（メモリは線形に増加しますが、トークンあたりの計算量は低く保たれます）

Cacheクラス

基本的なKVキャッシュインターフェースは、現在のトークンのキーとバリューのテンソルを受け取り、更新されたKテンソルとVテンソルを返します。これはモデルのforwardメソッドによって内部的に管理されます。

new_K, new_V = cache.update(k_t, v_t, layer_idx)
attn_output = attn_layer_idx_fn(q_t, new_K, new_V)

TransformersのCacheクラスを使用すると、自己注意モジュールは、過去の情報と現在の情報を統合するために、いくつかの重要なステップを実行します。

attention モジュールは、現在の kv ペアとキャッシュに保存されている過去の kv ペアを連結します。これにより、(new_tokens_length, past_kv_length + new_tokens_length) という形状の attention 重みが生成されます。現在の kv ペアと過去の kv ペアは本質的に結合され、attention スコアの計算に使用されることで、モデルが過去のコンテキストと現在の入力の両方を認識できるようになります。
forward メソッドが反復的に呼び出される場合、アテンションマスクの形状が、過去の kv ペアと現在の kv ペアの合計長と一致していることが重要です。アテンションマスクの形状は (batch_size, past_kv_length + new_tokens_length) である必要があります。これは通常 generate() 内で内部的に処理されますが、Cache を使用して独自の生成ループを実装する場合は、この点に注意してください。アテンションマスクには、過去のトークンと現在のトークンの値を保持する必要があります。

キャッシュストレージの実装

Cachesはレイヤーのリストとして構成されており、各レイヤーにはキーとバリューのキャッシュが含まれます。キーとバリューのキャッシュは、形状が[batch_size, num_heads, seq_len, head_dim]のテンソルです。

レイヤーには異なるタイプ（DynamicLayer、StaticLayer、StaticSlidingWindowLayer など）があり、主にシーケンス長の扱いとキャッシュの更新方法が異なります。

最もシンプルなものは、処理されるトークンが増えるにつれて成長するDynamicLayerです。シーケンス長の次元(seq_len)は、新しいトークンごとに増加します:

cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)

StaticLayer や StaticSlidingWindowLayer などの他のレイヤータイプは、キャッシュが作成される際に設定される固定のシーケンス長を持ちます。これにより、torch.compile との互換性が確保されます。StaticSlidingWindowLayer の場合、新しいトークンが追加されると、既存のトークンはキャッシュからシフトアウトされます。

以下の例は、DynamicCache を使用して生成ループを作成する方法を示しています。前述したように、アテンションマスクは、過去のトークン値と現在のトークン値の連結です。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from accelerate import Accelerator

device = Accelerator().device

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map=device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

past_key_values = DynamicCache(config=model.config)
messages = [{"role": "user", "content": "Hello, what's your name."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

generated_ids = inputs.input_ids
max_new_tokens = 10

for _ in range(max_new_tokens):
    outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
    # Greedily sample one next token
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    # Prepare inputs for the next generation step by leaving unprocessed tokens, in our case we have only one new token
    # and expanding attn mask for the new token, as explained above
    attention_mask = inputs["attention_mask"]
    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
    inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
"[INST] Hello, what's your name. [/INST]  Hello! My name is LLaMA,"

Update on GitHub