Cómo funciona la caché

Imagina que estás teniendo una conversación con alguien y, en lugar de recordar lo que dijo anteriormente, tiene que empezar desde cero cada vez que respondes. ¿No sería lento e ineficiente?

Puede extender esta analogía a los modelos transformer. La generación de modelos autoregresivos puede ser lenta porque realiza una predicción a la vez, un token cada vez. Cada nueva predicción depende de todo el contexto anterior.

Para predecir el token 1000, el modelo requiere información de los 999 tokens anteriores. La información se representa como multiplicaciones de matrices a través de las representaciones de los tokens.

Para predecir el token 1001, se necesita la misma información de los 999 tokens anteriores, además de cualquier información del token 1000. ¡Esto supone muchas multiplicaciones de matrices que un modelo debe calcular una y otra vez para cada token!

Una caché de clave-valor (KV) elimina esta ineficiencia almacenando pares kv derivados de las capas de atención de los tokens procesados previamente. Los pares kv almacenados se recuperan de la caché y se reutilizan para los tokens posteriores, evitando la necesidad de recalcularlos.

El uso de caché debe limitarse a la inferencia. Si se habilita durante el entrenamiento, puede provocar errores inesperados.

Para comprender mejor cómo y por qué funciona la caché, echemos un vistazo más detallado a la estructura de las matrices de atención.

Matrices de atención

La atención de producto escalar se calcula como se muestra a continuación para un lote de tamaño b, número de cabezas de atención h, longitud de secuencia hasta ahora T, y dimensión por cabeza de atención d_head. $$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$$

Las matrices de consulta (Q), clave (K) y valor (V) son proyecciones de las incrustaciones de entrada con forma (b, h, T, d_head).

Para la atención causal, la máscara impide que el modelo preste atención a tokens futuros. Una vez que un token se procesa, su representación nunca cambia con respecto a los tokens futuros, lo que significa que $K_{\text{past}}$ y $V_{\text{past}}$ pueden almacenarse en caché y reutilizarse para calcular la representación del último token. $$\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}])$$

En el momento de la inferencia, solo se necesita la consulta del último token para calcular la representación $x_t$ que predice el siguiente token $ t+1 $. En cada paso, los nuevos vectores de clave y valor se guardan en la caché y se añaden a las claves y valores anteriores. $$K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$$

La atención se calcula de forma independiente en cada capa del modelo, y la caché se realiza por capa.

Consulte la tabla siguiente para comparar cómo la caché mejora la eficiencia.

Clase de caché

Una interfaz básica de caché KV toma un tensor de clave y un tensor de valor para el token actual y devuelve los tensores K y V actualizados. Esto se gestiona internamente mediante el método forward de un modelo.

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

Cuando utiliza la clase Cache de Transformers, el módulo de autoatención realiza varios pasos críticos para integrar la información pasada y presente.

1. El módulo de atención concatena los pares kv actuales con los pares kv anteriores almacenados en la caché. Esto crea pesos de atención con la forma (new_tokens_length, past_kv_length + new_tokens_length). Los pares kv actuales y anteriores se combinan esencialmente para calcular las puntuaciones de atención, lo que garantiza que el modelo sea consciente del contexto previo y de la entrada actual.

2. Cuando se llama al método forward de forma iterativa, es crucial que la forma de la máscara de atención coincida con la longitud combinada de los pares kv pasados y actuales. La máscara de atención debe tener la forma (batch_size, past_kv_length + new_tokens_length). Esto normalmente se maneja internamente en generate(), pero si deseas implementar tu propio bucle de generación con Cache, ¡ten esto en cuenta! La máscara de atención debe contener los valores de los tokens pasados y actuales.

Implementación del almacenamiento en caché

Los cachés se estructuran como una lista de capas, donde cada capa contiene un caché de claves y un caché de valores. Los cachés de claves y valores son tensores con la forma [batch_size, num_heads, seq_len, head_dim].

Las capas pueden ser de diferentes tipos (p. ej. DynamicLayer, StaticLayer, StaticSlidingWindowLayer), lo que cambia principalmente cómo se maneja la longitud de la secuencia y cómo se actualiza la caché.

La más simple es una DynamicLayer que crece a medida que se procesan más tokens. La dimensión de longitud de secuencia (seq_len) aumenta con cada nuevo token:

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)

Otros tipos de capa, como StaticLayer y StaticSlidingWindowLayer, tienen una longitud de secuencia fija que se establece al crear la caché. Esto los hace compatibles con torch.compile. En el caso de StaticSlidingWindowLayer, los tokens existentes se desplazan fuera de la caché cuando se añade un token nuevo.

El siguiente ejemplo demuestra cómo crear un bucle de generación con DynamicCache. Como se discutió, la máscara de atención es una concatenación de los valores de los tokens pasados y actuales.

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