7 - Un modelo de lenguaje tipo GPT, desde cero

En el capítulo 7, Ana explica la estructura de un LLM tipo GPT: embeddings de los tokens con información de su posición, una pila de bloques Transformer con autoatención causal (cada token solo ve a sí mismo y a los tokens anteriores) y una proyección final, la cabeza LM, que convierte el último estado oculto en probabilidades para el siguiente token. Y explica el objetivo del entrenamiento, "sorprendentemente simple en concepto": predecir el siguiente token, con la entropía cruzada como pérdida, retropropagada por toda la red.

En este cuaderno construimos y entrenamos uno, diminuto, con PyTorch:

1. Importar las bibliotecas.
2. Un corpus en el que el contexto importa.
3. Del texto a los ejemplos de entrenamiento.
4. El modelo: atención causal, bloques Transformer y cabeza LM.
5. El entrenamiento: predecir el siguiente token.
6. Lo que aporta el contexto: el modelo frente a un trigrama.
7. Dónde mira el modelo: los pesos de atención.
8. Generar hojas de alta.

Paso 1: Importar las bibliotecas

---

PyTorch calcula automáticamente los gradientes de todos los parámetros, el paso que el libro llama retropropagar la pérdida "a través de toda la red". Las piezas del modelo, atención incluida, se escriben a mano.

import collections
import math
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn

torch.manual_seed(0)

<torch._C.Generator at 0x112e01350>

Paso 2: Un corpus en el que el contexto importa

---

Un LLM real se preentrena con una gran parte de internet. Nosotros generamos 3.000 hojas de alta sintéticas: un diagnóstico, algunos consejos generales, el fármaco que hay que tomar, más consejos y el signo de alarma por el que el paciente debería volver. El fármaco y el signo de alarma dependen del diagnóstico, que aparece varias frases antes; los consejos generales no dependen de nada. Para predecir bien el fármaco, un modelo tiene que mirar más atrás de las últimas palabras.

cases = {'neumonía': ('el antibiótico', 'dificultad para respirar'),
         'migraña': ('ibuprofeno', 'visión borrosa'),
         'hipertensión': ('las pastillas de la tensión', 'dolor en el pecho'),
         'gastroenteritis': ('suero oral', 'sangre en las heces'),
         'cistitis': ('fosfomicina', 'dolor de espalda')}
frequencies = ['dos veces al día', 'cada ocho horas', 'una vez al día']
general_advice = ['beba mucha agua .', 'descanse en casa unos días .', 'no conduzca hoy .', 'haga comidas ligeras .',
                  'evite el alcohol .', 'puede volver al trabajo la semana que viene .', 'guarde esta hoja .']
DIAGNOSIS, TAKE, WARNING = ['diagnóstico', ':'], ['tome'], ['vuelva', 'si', 'tiene']

def discharge_sheet(rng):
    """Una hoja de alta como lista de palabras: diagnóstico, algunos consejos, el fármaco, más consejos, el signo de alarma."""
    diagnosis = list(cases)[rng.integers(len(cases))]
    drug, warning = cases[diagnosis]
    words = DIAGNOSIS + [diagnosis, '.']
    for advice in rng.choice(general_advice, rng.integers(0, 4), replace=False):     # de 0 a 3 frases
        words += advice.split()
    words += TAKE + drug.split() + rng.choice(frequencies).split() + ['.']
    for advice in rng.choice(general_advice, rng.integers(0, 3), replace=False):     # de 0 a 2 frases
        words += advice.split()
    words += WARNING + warning.split() + ['.', '</s>']
    return words, diagnosis

rng = np.random.default_rng(0)
sheets = [discharge_sheet(rng) for _ in range(3000)]
vocab = sorted({w for words, _ in sheets for w in words} | {'<s>'})
token_id = {w: i for i, w in enumerate(vocab)}
stream = torch.tensor([token_id[w] for words, _ in sheets for w in ['<s>'] + words])   # todas las hojas, una detrás de otra

for words, _ in sheets[:2]:
    print(" ".join(words), "\n")
print(f"Vocabulario: {len(vocab)} tokens (uno por palabra), corpus: {len(stream):,} tokens, "
      f"hoja más larga: {max(len(words) for words, _ in sheets)} tokens")

diagnóstico : cistitis . descanse en casa unos días . haga comidas ligeras . tome fosfomicina dos veces al día . vuelva si tiene dolor de espalda . </s> 

diagnóstico : neumonía . tome el antibiótico una vez al día . guarde esta hoja . vuelva si tiene dificultad para respirar . </s> 

Vocabulario: 69 tokens (uno por palabra), corpus: 100,277 tokens, hoja más larga: 56 tokens

3.000 hojas forman un corpus de 100.277 tokens con un vocabulario de solo 69 palabras, un token por palabra (un LLM real usa tokens de subpalabras, como en el cuaderno Tokens y embeddings). La hoja más larga tiene 56 tokens, así que un contexto de 64 tokens siempre contiene una hoja entera.

Paso 3: Del texto a los ejemplos de entrenamiento

---

El ejemplo del libro: de "El médico diagnosticó al paciente con neumonía" salen los ejemplos Entrada = "El", Objetivo = "médico"; Entrada = "El médico", Objetivo = "diagnosticó"; y así sucesivamente. En la práctica, el texto se corta en ventanas de una longitud fija, aquí 64 tokens, y una ventana da todos esos ejemplos a la vez: la secuencia objetivo es la secuencia de entrada desplazada una posición a la izquierda. La máscara causal del paso siguiente se asegura de que la predicción en cada posición solo vea los tokens hasta esa posición.

context_length = 64

def random_windows(batch_size, generator):
    starts = torch.randint(0, len(stream) - context_length - 1, (batch_size,), generator=generator)
    x = torch.stack([stream[s:s + context_length] for s in starts])
    y = torch.stack([stream[s + 1:s + context_length + 1] for s in starts])     # la misma ventana, un token después
    return x, y

x, y = stream[None, :context_length], stream[None, 1:context_length + 1]    # la primera ventana del corpus
for t in range(6):
    print(f"Entrada = {' '.join(vocab[i] for i in x[0, :t + 1])!r:<45} Objetivo = {vocab[y[0, t]]!r}")

Entrada = '<s>'                                         Objetivo = 'diagnóstico'
Entrada = '<s> diagnóstico'                             Objetivo = ':'
Entrada = '<s> diagnóstico :'                           Objetivo = 'cistitis'
Entrada = '<s> diagnóstico : cistitis'                  Objetivo = '.'
Entrada = '<s> diagnóstico : cistitis .'                Objetivo = 'descanse'
Entrada = '<s> diagnóstico : cistitis . descanse'       Objetivo = 'en'

Paso 4: El modelo: atención causal, bloques Transformer y cabeza LM

---

Las mismas piezas que en el cuaderno La atención y el Transformer, ahora como capas entrenables de PyTorch:

• Autoatención causal con 4 cabezas: las consultas, las claves y los valores son proyecciones de los vectores de los tokens, $\text{softmax}\big(QK^\top / \sqrt{d_k}\big)V$, y una máscara triangular pone a $-\infty$ las puntuaciones de los tokens futuros antes de la softmax. Guardamos los últimos pesos de atención para mirarlos en el paso 7.
• Un bloque Transformer con Pre-LayerNorm, como en muchos LLM modernos: $x + \text{Atención}(\text{LN}(x))$ y después $x + \text{FFN}(\text{LN}(x))$, con conexiones residuales.
• Embeddings de los tokens más embeddings aprendidos de las posiciones, 2 bloques y la cabeza LM, $z_t = W_\text{vocab}\, h_t$, que comparte su matriz con los embeddings de los tokens (weight tying).

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.heads = heads
        self.qkv = nn.Linear(d_model, 3 * d_model)         # W_Q, W_K y W_V en una sola matriz
        self.out = nn.Linear(d_model, d_model)             # W_O
        self.register_buffer('mask', torch.tril(torch.ones(context_length, context_length)).bool())

    def forward(self, x):
        batch, T, d = x.shape
        q, k, v = self.qkv(x).split(d, dim=2)
        q, k, v = [m.view(batch, T, self.heads, d // self.heads).transpose(1, 2) for m in (q, k, v)]
        scores = q @ k.transpose(-2, -1) / math.sqrt(d // self.heads)
        scores = scores.masked_fill(~self.mask[:T, :T], float('-inf'))   # sin mirar los tokens futuros
        self.weights = scores.softmax(dim=-1)
        return self.out((self.weights @ v).transpose(1, 2).reshape(batch, T, d))

class TransformerBlock(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.attention = CausalSelfAttention(d_model, heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model))

    def forward(self, x):
        x = x + self.attention(self.norm1(x))              # Pre-LayerNorm y conexión residual
        return x + self.ffn(self.norm2(x))

class TinyGPT(nn.Module):
    def __init__(self, vocab_size, d_model=64, heads=4, layers=2):
        super().__init__()
        self.tokens = nn.Embedding(vocab_size, d_model)
        self.positions = nn.Embedding(context_length, d_model)
        self.blocks = nn.ModuleList([TransformerBlock(d_model, heads) for _ in range(layers)])
        self.norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.tokens.weight            # weight tying
        nn.init.normal_(self.tokens.weight, std=0.02)      # vectores iniciales pequeños, como en GPT
        nn.init.normal_(self.positions.weight, std=0.02)

    def forward(self, x):
        h = self.tokens(x) + self.positions(torch.arange(x.shape[1]))
        for block in self.blocks:
            h = block(h)
        return self.lm_head(self.norm(h))                  # logits del siguiente token en cada posición

model = TinyGPT(len(vocab))
print(f"Parámetros: {sum(p.numel() for p in model.parameters()):,}")

Parámetros: 108,608

108.608 parámetros. Un modelo del tamaño de GPT-3 tiene más de un millón de veces más, con las mismas piezas: más dimensiones, más cabezas y muchos más bloques.

Paso 5: El entrenamiento: predecir el siguiente token

---

1.500 pasos con lotes de 32 ventanas. En cada posición de cada ventana el modelo predice el siguiente token, y la pérdida es la entropía cruzada media de todas esas predicciones. loss.backward() calcula los gradientes de todos los parámetros, embeddings, proyecciones de la atención, capas feed-forward, normalizaciones y cabeza LM, y el optimizador, AdamW (una variante del descenso del gradiente), los actualiza todos a la vez.

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
generator = torch.Generator().manual_seed(0)
losses = []
for step in range(1, 1501):
    x, y = random_windows(32, generator)
    logits = model(x)
    loss = nn.functional.cross_entropy(logits.reshape(-1, len(vocab)), y.reshape(-1))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    losses.append(loss.item())
model.eval()
print(f"Pérdida en el primer paso: {losses[0]:.3f} (ln {len(vocab)} = {math.log(len(vocab)):.3f}), "
      f"media de los últimos 100 pasos: {np.mean(losses[-100:]):.3f}")

plt.figure(figsize=(8, 3.5))
plt.plot(losses, lw=0.8)
plt.xlabel('paso'); plt.ylabel('entropía cruzada'); plt.title('Entrenamiento prediciendo el siguiente token'); plt.grid(True); plt.show()

Pérdida en el primer paso: 4.219 (ln 69 = 4.234), media de los últimos 100 pasos: 0.326

En el primer paso, la pérdida es 4,219, casi exactamente $\ln 69$: el modelo sin entrenar da más o menos la misma probabilidad a cada uno de los 69 tokens. Al final está en torno a 0,326. No puede llegar a cero, porque parte de cada hoja es aleatoria, qué diagnóstico, qué consejos y en qué orden, y ningún modelo puede predecir eso. Lo que sí puede predecir casi exactamente son las palabras que se deducen del contexto.

Paso 6: Lo que aporta el contexto: el modelo frente a un trigrama

---

Generamos 500 hojas nuevas y, para cada una, damos al modelo el texto hasta la palabra anterior al fármaco, y después hasta la palabra anterior al signo de alarma, y dejamos que complete la frase de forma voraz, tomando siempre el token más probable, hasta el punto. Hacemos lo mismo con un modelo de trigramas, el n-grama del capítulo 8 con un contexto de dos palabras, contado con las mismas 3.000 hojas.

def next_token_probs(words):
    with torch.no_grad():
        return model(torch.tensor([[token_id[w] for w in words]]))[0, -1].softmax(dim=-1).numpy()

trigram_counts = collections.defaultdict(collections.Counter)
for words, _ in sheets:
    padded = ['<s>', '<s>'] + words
    for i in range(2, len(padded)):
        trigram_counts[tuple(padded[i - 2:i])][padded[i]] += 1

def complete(prefix, predict):
    words = list(prefix)
    while words[-1] != '.':
        words.append(predict(words))
    return " ".join(words[len(prefix):-1])

gpt_next = lambda words: vocab[next_token_probs(words).argmax()]
trigram_next = lambda words: trigram_counts[tuple((['<s>', '<s>'] + words)[-2:])].most_common(1)[0][0]

test_rng = np.random.default_rng(1000)
scores = {'GPT': [0, 0], 'trigrama': [0, 0]}
for _ in range(500):
    words, diagnosis = discharge_sheet(test_rng)
    drug, warning = cases[diagnosis]
    drug_prefix = ['<s>'] + words[:words.index(TAKE[-1]) + 1]
    warning_prefix = ['<s>'] + words[:len(words) - len(warning.split()) - 2]
    for name, predict in [('GPT', gpt_next), ('trigrama', trigram_next)]:
        scores[name][0] += complete(drug_prefix, predict).startswith(drug)
        scores[name][1] += complete(warning_prefix, predict) == warning
for name, (drug_ok, warning_ok) in scores.items():
    print(f"{name:<9} fármaco correcto: {drug_ok / 500:.1%}, signo de alarma correcto: {warning_ok / 500:.1%}")

example_rng = np.random.default_rng(7)
while True:                                          # un ejemplo con al menos 15 tokens entre el diagnóstico y el fármaco
    example, _ = discharge_sheet(example_rng)
    if example.index(TAKE[-1]) - len(DIAGNOSIS) >= 15:
        break
prefix = ['<s>'] + example[:example.index(TAKE[-1]) + 1]
p = next_token_probs(prefix)
print(f"\n{' '.join(prefix[1:])} ...")
print("GPT, siguiente token:", ", ".join(f"{vocab[i]} {p[i]:.3f}" for i in np.argsort(-p)[:3]))
counts = trigram_counts[tuple(prefix[-2:])]
print("Trigrama, siguiente token:", ", ".join(f"{w} {c / sum(counts.values()):.3f}" for w, c in counts.most_common(3)))

GPT       fármaco correcto: 100.0%, signo de alarma correcto: 100.0%
trigrama  fármaco correcto: 22.2%, signo de alarma correcto: 22.2%

diagnóstico : hipertensión . no conduzca hoy . descanse en casa unos días . guarde esta hoja . tome ...
GPT, siguiente token: las 0.999, borrosa 0.000, casa 0.000
Trigrama, siguiente token: las 0.208, el 0.202, fosfomicina 0.201

El modelo GPT completa el fármaco correcto y el signo de alarma correcto en el 100,0% de las 500 hojas nuevas. El trigrama acierta el 22,2% de cada uno: solo ve las dos últimas palabras, las mismas para todos los diagnósticos, así que siempre completa la misma frase y solo acierta cuando el diagnóstico resulta ser el que le corresponde. En el ejemplo, con tres frases de consejos entre el diagnóstico y el fármaco, el modelo da 0,999 a "las" (de "las pastillas de la tensión"), mientras que el trigrama reparte su probabilidad entre los comienzos de todos los fármacos.

Paso 7: Dónde mira el modelo: los pesos de atención

---

Para predecir el fármaco, el modelo tiene que encontrar el diagnóstico. Pasamos la hoja de ejemplo hasta la palabra anterior al fármaco y mostramos, para cada una de las 4 cabezas de los 2 bloques, los pesos de atención desde esa última posición a todos los tokens anteriores.

with torch.no_grad():
    model(torch.tensor([[token_id[w] for w in prefix]]))
diagnosis_position = prefix.index(DIAGNOSIS[-1]) + 1
rows, labels = [], []
for b, block in enumerate(model.blocks, start=1):
    weights = block.attention.weights[0, :, -1].numpy()             # (cabezas, posiciones): desde el último token
    for head, w in enumerate(weights, start=1):
        rows.append(w); labels.append(f'bloque {b}, cabeza {head}')
        print(f"bloque {b}, cabeza {head}: peso en el diagnóstico {w[diagnosis_position]:.3f}, "
              f"el mayor en {prefix[w.argmax()]!r} ({w.max():.3f})")

plt.figure(figsize=(14, 4))
plt.imshow(np.array(rows), cmap='Blues', vmin=0, vmax=1, aspect='auto')
plt.xticks(range(len(prefix)), prefix, rotation=90); plt.yticks(range(len(labels)), labels)
plt.gca().get_xticklabels()[diagnosis_position].set_color('red')
plt.colorbar(label='peso de atención'); plt.title(f'Atención desde {prefix[-1]!r} a los tokens anteriores (el diagnóstico en rojo)')
plt.tight_layout(); plt.show()

bloque 1, cabeza 1: peso en el diagnóstico 0.008, el mayor en 'tome' (0.164)
bloque 1, cabeza 2: peso en el diagnóstico 0.926, el mayor en 'hipertensión' (0.926)
bloque 1, cabeza 3: peso en el diagnóstico 0.027, el mayor en 'tome' (0.334)
bloque 1, cabeza 4: peso en el diagnóstico 0.013, el mayor en 'tome' (0.276)
bloque 2, cabeza 1: peso en el diagnóstico 0.028, el mayor en 'tome' (0.436)
bloque 2, cabeza 2: peso en el diagnóstico 0.001, el mayor en 'tome' (0.902)
bloque 2, cabeza 3: peso en el diagnóstico 0.029, el mayor en ':' (0.350)
bloque 2, cabeza 4: peso en el diagnóstico 0.001, el mayor en 'tome' (0.950)

Nadie le ha dicho al modelo dónde está el diagnóstico. Aun así, en el primer bloque, la cabeza 2 pone 0,926 de su atención en "hipertensión", y casi nada en los consejos intermedios. Las demás cabezas y el segundo bloque miran sobre todo tokens cercanos, como "tome" o los dos puntos. La pérdida del entrenamiento por sí sola, predecir el siguiente token, ha hecho que una cabeza se especialice en encontrar el diagnóstico.

Paso 8: Generar hojas de alta

---

Para generar, el modelo lee el texto hasta el momento, convierte los logits de la última posición en probabilidades con una softmax a temperatura $T$, muestrea el siguiente token, lo añade y repite, token a token, hasta </s>. Generamos 200 hojas con $T = 1$ y comprobamos si cada una es coherente: el fármaco y el signo de alarma son los que corresponden a su diagnóstico.

def generate(sample_rng, temperature=1.0):
    words = ['<s>']
    while words[-1] != '</s>' and len(words) < context_length:
        with torch.no_grad():
            logits = model(torch.tensor([[token_id[w] for w in words]]))[0, -1].numpy() / temperature
        p = np.exp(logits - logits.max())
        words.append(vocab[sample_rng.choice(len(vocab), p=p / p.sum())])
    return words[1:]

def coherent(words):
    text = " ".join(words)
    return any(text.startswith(" ".join(DIAGNOSIS + [d])) and " ".join(TAKE + drug.split()) in text
               and " ".join(WARNING + warning.split() + ['.', '</s>']) in text for d, (drug, warning) in cases.items())

sample_rng = np.random.default_rng(42)
generated = [generate(sample_rng) for _ in range(200)]
for words in generated[:3]:
    print(" ".join(words), "\n")
print(f"Hojas coherentes: {sum(coherent(w) for w in generated)} de {len(generated)}")
print("La primera incoherente:", " ".join(next(w for w in generated if not coherent(w))))
print(f"Hojas idénticas a alguna del corpus: {sum(tuple(w) in {tuple(s) for s, _ in sheets} for w in generated)} de {len(generated)}")

diagnóstico : neumonía . beba mucha agua . evite el alcohol . tome el antibiótico dos veces al día . vuelva si tiene dificultad para respirar . </s> 

diagnóstico : migraña . guarde esta hoja . beba mucha agua . no conduzca hoy . tome ibuprofeno dos veces al día . beba mucha agua . vuelva si tiene visión borrosa . </s> 

diagnóstico : migraña . no conduzca hoy . descanse en casa unos días . guarde esta hoja . tome ibuprofeno una vez al día . haga comidas ligeras . vuelva si tiene visión borrosa . </s> 

Hojas coherentes: 196 de 200
La primera incoherente: diagnóstico : neumonía : hipertensión . tome las pastillas de la tensión una vez al día . guarde esta hoja . vuelva si tiene dolor en el pecho . </s>

Hojas idénticas a alguna del corpus: 59 de 200

196 de las 200 hojas generadas son coherentes. Los consejos cambian de una hoja a otra, porque se muestrean, y 59 de las 200 resultan idénticas a alguna hoja del corpus: con tan pocas combinaciones posibles, el modelo cae a menudo en una que ya ha visto, pero la mayoría de las hojas son nuevas. La primera incoherente mezcla dos diagnósticos, neumonía e hipertensión, y después sigue con el tratamiento de la hipertensión. A la escala de este juguete, es el mismo fenómeno que el libro llama alucinación: el modelo genera lo probable, no lo verdadero.

Conclusión

---

• Un modelo tipo GPT es una pila de bloques Transformer con atención causal entre los embeddings de los tokens y una cabeza LM que da las probabilidades del siguiente token.
• Se entrena con un único objetivo, predecir el siguiente token, y una única pérdida, la entropía cruzada, cuyo gradiente llega a la vez a todos los parámetros de la red.
• La atención le permite usar información muy anterior del texto, algo que un n-grama no puede hacer: es la diferencia entre predecir casi siempre el fármaco correcto y adivinar.
• Generar texto es repetir la predicción token a token, muestreando de las probabilidades.
• Nuestro modelo tiene unos cien mil parámetros y 3.000 hojas muy regulares; los LLM del libro tienen miles de millones de parámetros y una gran parte de internet, pero el objetivo y la arquitectura son los mismos. El cuaderno Preentrenamiento y fine-tuning da el siguiente paso del libro: adaptar un modelo preentrenado a una tarea nueva.