# キャッシュの仕組み

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

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

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

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

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

> [!WARNING]
> **推論**での使用に限定してください。学習中に有効にすると、予期しないエラーが発生する可能性があります。

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

## アテンション行列

**スケーリングドットプロダクトアテンション**は、バッチサイズ`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}]) $$

推論時において、次のトークン $ t+1 $ を予測する表現 $ 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`メソッドによって内部的に管理されます。

```py
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](/docs/transformers/main/ja/internal/generation_utils#transformers.Cache) クラスを使用すると、自己アテンションモジュールは、過去の情報と現在の情報を統合するために、いくつかの重要なステップを実行します。

1. attention モジュールは、現在の kv ペアとキャッシュに保存されている過去の kv ペアを連結します。これにより、`(new_tokens_length, past_kv_length + new_tokens_length)` という形状の attention 重みが生成されます。現在の kv ペアと過去の kv ペアは本質的に結合され、attention スコアの計算に使用されることで、モデルが過去のコンテキストと現在の入力の両方を認識できるようになります。

2. `forward` メソッドが反復的に呼び出される場合、アテンションマスクの形状が、過去の kv ペアと現在の kv ペアの合計長と一致していることが重要です。アテンションマスクの形状は `(batch_size, past_kv_length + new_tokens_length)` である必要があります。これは通常 [generate()](/docs/transformers/main/ja/main_classes/text_generation#transformers.GenerationMixin.generate) 内で内部的に処理されますが、[Cache](/docs/transformers/main/ja/internal/generation_utils#transformers.Cache) を用いて独自の生成ループを実装する場合は、この点に注意してください。アテンションマスクには、過去のトークンと現在のトークンの値を保持する必要があります。

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

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

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

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

```py
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](/docs/transformers/main/ja/internal/generation_utils#transformers.DynamicCache) を使用して生成ループを作成する方法を示しています。前述したように、アテンションマスクは、過去のトークン値と現在のトークン値の連結です。

```py
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,"
```