Generación: decoding, sampling y KV cache

Tu modelo entrenado sabe una sola cosa: dado un prefijo, producir una distribución sobre el siguiente token. Generar texto es aplicar eso en bucle: elegir un token, agregarlo al contexto y repetir. La parte interesante está en cómo se elige, y en cómo evitar recalcular todo en cada paso.

El bucle autoregresivo

import torch
import torch.nn.functional as F

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None):
    model.eval()
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -model.cfg.block_size:]       # recorta al contexto máximo
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :]                        # solo la última posición

        if temperature == 0:                             # greedy
            next_id = logits.argmax(dim=-1, keepdim=True)
        else:
            logits = logits / temperature
            if top_k is not None:
                logits = top_k_filter(logits, top_k)
            if top_p is not None:
                logits = top_p_filter(logits, top_p)
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)

        idx = torch.cat([idx, next_id], dim=1)
    return idx

Fíjate en idx[:, -block_size:]: como los embeddings posicionales son aprendidos, el modelo no puede ver más de block_size tokens. Cuando el texto crece, se descarta lo más antiguo.

Greedy: siempre el más probable

Con temperature=0 elegimos $\arg\max$. Es determinista y suele dar buenos resultados en respuestas cortas, pero en textos largos tiende a caer en bucles repetitivos: una vez que el modelo repite una frase, la frase repetida se vuelve todavía más probable.

Temperatura

Para muestrear, dividimos los logits $z$ por una temperatura $\tau > 0$ antes del softmax:

la probabilidad de cada token sale de su puntaje dividido por la temperatura $\tau$: con $\tau$ chica casi siempre gana el favorito; con $\tau$ grande, la suerte se reparte más

• $\tau < 1$ afila la distribución: más conservador, más cerca de greedy.
• $\tau = 1$ usa la distribución que aprendió el modelo, tal cual.
• $\tau > 1$ la aplana: más diverso, más errores.

Mueve el slider y observa cómo cambia la masa de probabilidad:

Top-k y top-p: cortar la cola

El problema de muestrear directamente es la cola larga: miles de tokens con probabilidad individual minúscula pero que, sumados, tienen una probabilidad no despreciable. Cada tanto se elige uno de ellos y el texto se descarrila.

Top-k se queda con los $k$ tokens más probables y descarta el resto:

def top_k_filter(logits, k):
    v, _ = torch.topk(logits, min(k, logits.size(-1)))
    return logits.masked_fill(logits < v[:, [-1]], float("-inf"))

El problema de top-k es que $k$ es fijo: cuando el modelo está muy seguro (solo 2 tokens plausibles) deja pasar basura; cuando hay muchas opciones válidas (el comienzo de una oración) corta de más.

Top-p o nucleus sampling (Holtzman et al., 2020) se adapta: se queda con el conjunto más chico de tokens cuya probabilidad acumulada alcanza $p$:

ordenar los tokens de más a menos probable e ir sumando hasta llegar a $p$: se sortea solo entre esos

def top_p_filter(logits, p):
    sorted_logits, sorted_idx = torch.sort(logits, descending=True)
    probs = F.softmax(sorted_logits, dim=-1)
    cum = probs.cumsum(dim=-1)
    # descarta un token si la masa ANTERIOR a él ya alcanzaba p
    remove = (cum - probs) >= p
    sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
    return torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits)

Restar probs a la suma acumulada asegura que el token que cruza el umbral se conserve, y que el más probable nunca se descarte.

El costo oculto: recalcular el pasado

Mira el bucle de generate: en el paso $t$ pasamos los $t$ tokens completos por el modelo para usar solo el último logit. Pero las keys y values de los tokens anteriores no cambiaron: gracias a la máscara causal, la representación de la posición $j$ solo depende de los tokens $\le j$. Recalcularlas es trabajo tirado.

Contemos para generar $N$ tokens (con $d$ = n_embd):

El KV cache guarda, por capa, las keys y values ya calculadas. Cada paso nuevo procesa un único token: calcula su $q$, $k$, $v$, agrega $k$ y $v$ al cache y atiende sobre todo el cache.

Implementación

Tres cambios pequeños. En la atención, aceptamos un cache (un dict por capa) y recortamos la máscara para queries que empiezan en la posición Tk - T:

# dentro de CausalSelfAttention
def forward(self, x, cache=None):
    B, T, C = x.shape
    hs = C // self.n_head
    q, k, v = self.qkv(x).split(C, dim=2)
    q = q.view(B, T, self.n_head, hs).transpose(1, 2)
    k = k.view(B, T, self.n_head, hs).transpose(1, 2)
    v = v.view(B, T, self.n_head, hs).transpose(1, 2)

    if cache is not None:
        if "k" in cache:
            k = torch.cat([cache["k"], k], dim=2)
            v = torch.cat([cache["v"], v], dim=2)
        cache["k"], cache["v"] = k, v

    Tk = k.size(2)
    att = (q @ k.transpose(-2, -1)) / math.sqrt(hs)
    att = att.masked_fill(~self.mask[:, :, Tk - T : Tk, :Tk], float("-inf"))
    att = self.attn_drop(F.softmax(att, dim=-1))
    y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
    return self.resid_drop(self.proj(y))

En Block, pasamos el cache: x = x + self.attn(self.ln1(x), cache). En GPT.forward, agregamos caches y start_pos para que los embeddings posicionales sigan siendo correctos:

def forward(self, idx, targets=None, caches=None, start_pos=0):
    B, T = idx.shape
    assert start_pos + T <= self.cfg.block_size
    pos = torch.arange(start_pos, start_pos + T, device=idx.device)
    x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
    for i, block in enumerate(self.blocks):
        x = block(x, None if caches is None else caches[i])
    logits = self.lm_head(self.ln_f(x))
    loss = None
    if targets is not None:
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
    return logits, loss

La generación queda en dos fases: prefill (procesar el prompt completo de una vez) y decode (un token por paso):

@torch.no_grad()
def generate_cached(model, idx, max_new_tokens, temperature=1.0):
    model.eval()
    assert idx.size(1) + max_new_tokens <= model.cfg.block_size
    caches = [{} for _ in model.blocks]
    logits, _ = model(idx, caches=caches, start_pos=0)          # prefill
    pos = idx.size(1)
    for _ in range(max_new_tokens):
        probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, next_id], dim=1)
        logits, _ = model(next_id, caches=caches, start_pos=pos)  # decode
        pos += 1
    return idx

Verifica la equivalencia: con la misma semilla, generate y generate_cached deben producir la misma secuencia (salvo diferencias numéricas mínimas en GPU). La limitación de esta versión: con embeddings posicionales aprendidos no podemos deslizar la ventana, así que el total queda acotado a block_size.

El precio: memoria

El cache guarda $K$ y $V$ para cada capa, cada token y toda la dimensión del modelo:

K y V (por eso el 2), en cada una de las $L$ capas, para cada uno de los $T$ tokens, con $d$ números cada uno

Para GPT-2 small en fp16: $2 \cdot 12 \cdot 768 \cdot 2 = 36{,}864$ bytes por token, unos 37,7 MB para una secuencia de 1024 tokens. Multiplica por el tamaño del batch y por modelos con decenas de capas y contextos de cientos de miles de tokens: en inferencia de LLM grandes, el KV cache suele ser el cuello de botella de memoria, y técnicas como multi-query o grouped-query attention existen precisamente para achicarlo compartiendo keys y values entre cabezas.

Lo que te llevas

• Generar = bucle de "predecir, elegir, agregar".
• Greedy es determinista pero repetitivo; temperatura, top-k y top-p regulan el riesgo.
• El KV cache evita recalcular keys y values: pasa la atención total de $O(N^3 d)$ a $O(N^2 d)$.
• Su costo es memoria lineal en el largo del contexto, que domina la inferencia a gran escala.

Con esto cerraste el ciclo completo: texto → tokens → vectores → atención → bloques → entrenamiento → texto nuevo. Ya construiste un GPT.