Entrenar un GPT pequeño con cuentos populares indios: tokenizador, máscara y pérdida.
Un GPT de 9M parámetros creado desde cero: por qué el tokenizador de GPT-2 falla con palabras de cuentos populares, cómo funcionan BPE y la máscara causal, y cómo un solo lote genera 4,096 pérdidas de token siguiente.
Fuentes: Building LLM, Part 1 — Training: Predicting the Next Token
Un pequeño GPT entrenado con cuentos populares indios
Cualquiera que haya creado un Vision Transformer ya conoce aproximadamente el 70 por ciento de los componentes de un modelo de lenguaje. Los embeddings, la atención, la pila residual y el bucle de entrenamiento son los mismos en ambos. Lo que queda es la diferencia, y esta guía lo aborda al construir un pequeño GPT desde cero y entrenarlo directamente en la GPU de una laptop. El corpus utilizado son cuentos populares indios: historias Jataka, fábulas al estilo Panchatantra y relatos de Punjab y Bengala.
Un verdadero problema surge desde el principio. El tokenizador original de GPT-2 divide palabras como Brahman y Rajah en fragmentos porque nunca las había encontrado, por lo que se entrena un tokenizador a partir del corpus. Al final se obtiene un modelo de 9 millones de parámetros que ha pasado de ser ruido aleatorio a tener pesos entrenados, y cualquiera puede reproducir esos resultados ejecutando el proceso de entrenamiento, que dura aproximadamente diez minutos.
Por qué la visión vino primero
Analizar imágenes antes que texto puede parecer el orden incorrecto, ya que los chatbots trabajan con palabras. Sin embargo, un transformador no se preocupa por si su entrada son fragmentos de fotos o partes de oraciones. Las cabezas de atención, LayerNorm, conexiones saltadas, MLP y el bucle de optimización no cambian. La pregunta más relevante es dónde resulta más sencillo aprender ese motor común, y las imágenes ganan por cuatro razones.
- Un token es simplemente un fragmento de imagen, una cuadrícula fija extraída de la foto; por lo tanto, no existe vocabulario, ni tokenizador, ni problema de palabras desconocidas que resolver primero.
- No hay máscara causal. Un ViT ve toda la imagen de una sola vez, por lo que la atención simple precede a las reglas adicionales del lenguaje que prohíben echar un vistazo al futuro.
- El comportamiento es visible. Los fragmentos, los mapas de atención y las regiones en las que se enfoca el modelo pueden dibujarse, y los píxeles resultan tangibles en comparación con el texto.
- Se entrena desde cero en una sola GPU en minutos (en CIFAR-10); así, cada cifra es un número que usted mismo genera y no se toma por algo dado de antemano.
En resumen, una vez que se construye un ViT, ya existe el transformador completo, solo adaptado para el procesamiento de imágenes. Un modelo de lenguaje únicamente añade cómo el texto se convierte en tokens (una búsqueda y no un cálculo), la máscara causal y la predicción del siguiente token. La parte difícil y abstracta, el lenguaje en sí, es mucho más fácil de comprender una vez que la máquina ya está familiarizada con él.
El entrenamiento es sencillo: adivinar el siguiente token
Durante el entrenamiento, todo lo que hace un modelo de lenguaje se reduce a leer los tokens hasta el momento y predecir el que sigue.
La figura ilustra cómo se predice el siguiente token a partir de los tokens anteriores. Si se le proporciona la frase inicial “En un cierto pueblo vivía un pobre brahmán”, se obtiene una probabilidad para cada candidato que podría seguir. Un punto, una coma o una palabra como quién o llamado serían todos gramaticalmente válidos. Nadie necesita anotar nada: la respuesta correcta es el token que aparecía a continuación en el texto. Por eso, solo el texto en bruto puede enseñar a un modelo de lenguaje, y constituye toda la señal de aprendizaje. Cada otro componente existe para respaldar esta única predicción, y ViT ya proporciona la mayoría de ellos.
En código, esta idea se resume en dos líneas: ejecutar el modelo y luego normalizar sus puntuaciones para obtener una distribución sobre el siguiente token.
logits, _ = model(ids) # ids = the tokens so far, shape (1, T)
probs = F.softmax(logits[0, -1], dim=-1) # one probability per vocabulary token
El resto explica cómo se construye y entrena model para que esas probabilidades resulten adecuadas.
Por qué es importante el tokenizador
Un ViT divide una imagen en parches, mientras que un modelo de lenguaje divide el texto en tokens. Un parche se calcula a partir de los valores de píxeles, pero un token de texto se obtiene de un vocabulario fijo. Alguien debe crear ese vocabulario con antelación, y quien lo elabora decide qué elementos el modelo considera fáciles.
El primer intento reutilizó r50k_base, el tokenizador listo para usar de GPT-2, cuyo vocabulario estaba basado en textos web occidentales de alrededor de 2019. La siguiente figura muestra cómo trata las palabras de los cuentos.
Brahman se divide en Brah y man, Rajah en Raj y ah, y Panchatantra se fragmenta en cinco partes. Ninguno de estos términos es poco conocido. Rajah aparece 709 veces en el corpus, pero el tokenizador nunca encontró estas palabras con suficiente frecuencia en sus propios datos de entrenamiento como para asignarle un token a cada una. La fragmentación alarga las secuencias y obliga al modelo a reconstituir Brahman a partir de Brah más man, donde man también se comparte con woman y human. Con solo alrededor de un millón de tokens de datos de entrenamiento, eso representa una verdadera desventaja.
La solución es un tokenizador entrenado en esos relatos, utilizando la codificación de pares de bytes (BPE). GPT emplea el mismo algoritmo, solo adaptado a un corpus más adecuado. El material de partida son doce colecciones de cuentos populares indios de dominio público del Proyecto Gutenberg, que incluyen dos libros de Jataka, dos volúmenes de cuentos infantiles indios, Old Deccan Days, Folk-Tales of Bengal, The Talking Thrush, Tales of the Punjab, Twenty-Two Goblins, dos colecciones de cuentos santal y Tales of the Sun from Southern India, sumando un total de aproximadamente 809,000 palabras.
La figura muestra los procesos de fusión de pares de bytes utilizados para entrenar el tokenizador. BPE comienza con los 256 valores de byte brutos, de modo que no puede haber nada desconocido, y luego fusiona repetidamente la pareja vecina más común para crear un nuevo token. Tras 7,935 fusiones en los relatos, el vocabulario cuenta con 8,192 entradas e incluye unidades completas como Rajah, Brahman y jackal.
Hay una decisión de diseño que merece una aclaración clara: el tokenizador de GPT-2 se reemplaza por completo, no se extiende.
La figura muestra el vocabulario personalizado de 8,192 tokens. Añadir tokens a los 50,257 existentes en r50k es una opción razonable cuando se desea mantener las capacidades en inglés de un modelo preentrenado grande, como lo hacen proyectos como OpenHathi. Entrenar desde cero con un corpus pequeño y especializado es otra situación, en la que el vocabulario de r50k basado en la web occidental solo ocupa espacio, y los relatos abarcan solo el 37 por ciento de él. El nuevo vocabulario de 8k es más pequeño, se utiliza en el 96.6 por ciento de los casos y está adaptado a los relatos, por lo que el reemplazo es la opción más adecuada.
Entrenar el tokenizador requiere unas pocas líneas de código con la biblioteca tokenizers de HuggingFace. El script lee cada libro de texto en formato raw, elimina los encabezados y pies de página de Project Gutenberg y entrena un modelo BPE.
import glob
from tokenizers import Tokenizer, models, trainers, pre_tokenizers
from llm_transformer.prep_tales import strip_gutenberg, RAW_DIR
# read every raw .txt book and strip its Project Gutenberg header/footer,
# leaving just the story text - one big string per book
tales_texts = [strip_gutenberg(open(p, encoding="utf-8", errors="ignore").read())
for p in sorted(glob.glob(f"{RAW_DIR}/*.txt"))] # 12 books, ~809k words
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) # start from bytes
trainer = trainers.BpeTrainer(vocab_size=8192, special_tokens=["<|endoftext|>"],
initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator(tales_texts, trainer) # learn the 7,935 merges
tok.save("tales_bpe.json")
El archivo de salida .json es el tokenizador final, una lista fija de fusiones. No se entrena nuevamente. Se carga y se utiliza para convertir texto en identificadores y viceversa. Una sola llamada a .encode() vincula el texto en bruto con la red neuronal.
tok = Tokenizer.from_file("tales_bpe.json") # reload the trained tokenizer
ids = tok.encode("Once upon a time").ids # -> [412, 987, 15, 733] (integers)
idx = torch.tensor(ids).unsqueeze(0) # (1, 4) — a batch of one sequence
vocab_size = tok.get_vocab_size() # 8192 - sets the NUMBER OF ROWS below
Solo esos enteros llegan a la red, y las palabras originales ya no existen para ella. Cada id es simplemente una posición en una tabla, y vocab_size determina cuántas entradas reserva la capa de embedding.
De ids a vectores
Una vez que el texto se convierte en una lista de ids, cada id se transforma en un vector. Este es el único punto en el que un LLM es más sencillo que un ViT.
La figura relaciona un índice del tokenizador con un vector de embedding. En un ViT, una pequeña multiplicación matricial calcula el token de cada parche a partir de píxeles. En un LLM, el token es un id discreto, por lo que el vector es simplemente el número de fila del id en una tabla aprendida, una búsqueda directa. El token 789 recupera la fila 789, que es un vector de 256 dimensiones; al añadirle un embedding posicional, como en el ViT, se completa la etapa de entrada.
En el código, una capa nn.Embedding funciona como la tabla, y al llamarla se realiza la búsqueda mediante indexación, no por multiplicación. Su .weight es un tensores aprendible de 8192 por 256 que contiene 2,1 millones de pesos entrenados mediante retropropagación, y una fila determinada solo cambia en los pasos en los que aparece su token.
self.token_embed = nn.Embedding(vocab_size, n_embed) # the 8,192 x 256 table
# .weight IS a learnable (8192, 256) parameter tensor — 2.1M weights, trained by backprop;
# each id's row is nudged only on steps where that token appears in the batch
...
tok_emb = self.token_embed(idx) # idx (B, T) -> vectors (B, T, 256), one row per id
Todo lo que viene después de este punto, es decir, la atención, el flujo residual y los bloques MLP, es idéntico al Vision Transformer, con una excepción.
La máscara causal
Un ViT permite que cada parche preste atención a todos los demás. Un modelo de lenguaje no puede hacerlo, porque su función es adivinar qué viene a continuación. Si la posición t viera la posición t+1, simplemente copiaría la respuesta. La solución es enmascarar la atención para que cada token solo vea a sí mismo y a los tokens anteriores.
La máscara que aparece en la figura es un triángulo inferior. Antes de aplicar softmax, todas las posiciones futuras reciben un valor de infinito negativo, que softmax convierte en un peso de cero exacto, de modo que los tokens bloqueados no aportan nada. Este triángulo representa la única diferencia arquitectónica entre la atención de ViT y la atención de los LLM. Los productos escalar, la escala mediante la raíz cuadrada del tamaño de la cabeza, softmax y la suma ponderada de los valores siguen siendo los mismos.
En el código, se trata de una línea adicional encima de la atención habitual. Se crea una matriz triangular inferior compuesta por unos, se calculan las puntuaciones como de costumbre, y las posiciones con valor cero se rellenan con infinito negativo antes de aplicar softmax.
# tril = torch.tril(torch.ones(T, T)) — a lower-triangular matrix of 1s, made once
attn = q @ k.transpose(-2, -1) / math.sqrt(head_size) # scores, (B, T, T)
attn = attn.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # blank out the future
attn = F.softmax(attn, dim=-1) # -inf -> weight 0
out = attn @ v
Una sola pasada entrena todas las longitudes de contexto
La máscara de una sola línea oculta algo que merece una mirada más detallada. La gente tiende a imaginar que un modelo de lenguaje funciona como si generara texto: añadiendo un token al contexto en constante crecimiento y volviendo a calcular la atención en cada paso. Eso es cierto durante la inferencia, pero no durante el entrenamiento. En el entrenamiento, la oración se conoce por completo, por lo que Q, K y V para todos los tokens T se calculan juntos, formando toda la matriz de atención T por T en un solo paso.
Esa única pasada solo es válida gracias a la máscara causal. Considere la fila t de la matriz enmascarada: todas las entradas a la derecha de la diagonal son infinito negativo, por lo que se convierten en cero después del softmax. Al calcular out = attn @ V, la fila t recibe una suma ponderada solo de los primeros t vectores de valores. Por lo tanto, todas las longitudes de contexto se procesan al mismo tiempo.
La figura muestra diferentes ejecuciones y la forma de la máscara durante el entrenamiento. Úsela fila por fila:
- La fila 1 solo tiene en cuenta el token 1, por lo que su salida predice el token 2 a partir de un contexto de un solo token.
- La fila 3 tiene en cuenta los tokens del 1 al 3 y predice el token 4 a partir de tres tokens de contexto.
- La fila 6 tiene en cuenta los seis tokens y predice el token 7 a partir del contexto completo.
Aquí hay otra perspectiva. El contexto visible de cada fila forma una matriz triangular inferior que crece en uno cada vez: 1 por 1 para el primer token, luego 2 por 2 y 3 por 3. Estos son los bloques anidados en la esquina superior izquierda de la matriz grande, dibujados en naranja; así, una sola pasada de 6 por 6 contiene todas las secuencias más pequeñas, y la franja en la parte inferior las despliega. Todas se generan dentro de un attn @ V, y no hay bucles que recorran posiciones en Python. Esto explica la eficiencia en el entrenamiento de los transformadores: una secuencia con T tokens equivale a T ejemplos de entrenamiento, con contextos que van desde la longitud 1 hasta T, evaluados en paralelo. La misma máscara triangular que proporciona esos ejemplos también impide que cualquier posición mire hacia adelante durante la generación, la cual se realiza fila por fila y mantiene cada paso económico al almacenar en caché las claves y valores anteriores (el KV-cache).
Números reales a través de la atención
Hasta ahora, cada celda solo indicaba asistir o no asistir. Al mostrar los valores multiplicados, se presentan las cifras reales para las mismas tres ejecuciones, extraídas del bloque 0, cabeza 0 del modelo entrenado en la oración “Once there lived a clever jackal”.
La figura muestra los cálculos de Q, K y V a lo largo de las ejecuciones. Dentro de una misma fila, los pesos de atención son positivos y suman 1, por lo que la salida es una combinación convexa de los vectores de valor visibles, que corresponde a un promedio ponderado.
- En la ejecución 1, hay un token con peso 1.00, por lo que
out1 = V1. La primera posición copia su propio valor, ya que aún no hay nada con lo mezclar. - En la ejecución 3, la salida combina tres valores utilizando pesos de aproximadamente 0.20 para el primer token, 0.06 para el segundo y 0.74 para el tercero. Se basa principalmente en su propio token, lived, pero también incorpora una quinta parte de Once; por lo tanto, se trata de una verdadera mezcla.
Eso abarca todo lo que hace la atención: las puntuaciones de las consultas y las claves establecen los pesos, y estos pesos mezclan los valores. El enmascaramiento garantiza que la mezcla nunca trascienda el token actual, por lo que cada fila representa un resumen equitativo de todo lo leído hasta el momento. Las cabezas reales tienen 32 dimensiones de valor, y el modelo cuenta con ocho cabezas en cada una de las seis capas. Las franjas muestran las primeras 8 dimensiones de una sola cabeza, aunque el mismo mecanismo se aplica en todas partes.
Predicciones T, pérdidas T
La predicción del token siguiente aporta una ventaja que el ViT carecía. El ViT realizaba una sola predicción por imagen, un etiqueta de clase. La máscara causal permite a un modelo de lenguaje adivinar el token siguiente en todas las posiciones al mismo tiempo, de modo que una secuencia genera tantas señales de aprendizaje como tokens tenga.
La figura muestra el promedio de T predicciones y valores de pérdida para una entrada de T tokens en cada ejecución. Cada posición recibe una pérdida por entropía cruzada, que mide cuán sorprendido estaba el modelo por el verdadero token siguiente. Algunas son bajas, como a después de lived, con solo 0.58. Otras son altas, como el adjetivo clever, con 7.17, ya que podrían encajar muchas palabras diferentes. El promedio de las seis da 3.59, y la retropropagación disminuye ese valor utilizando el mismo mecanismo de regla de cadena que el ViT.
En el código, todas las pérdidas T provienen de una única llamada a cross_entropy. La matriz de posiciones (B, T) se aplana en una sola lista larga, y cada elemento se evalúa en comparación con su siguiente token real.
logits = self.lm_head(x) # (B, T, vocab)
B, T, C = logits.shape
loss = F.cross_entropy(logits.reshape(B*T, C), # every position ...
targets.reshape(B*T)) # ... vs the token that actually followed
Cada pérdida es una clasificación de vocabulario
Es importante ser preciso respecto a lo que compara cross_entropy, ya que tanto la predicción como el objetivo son vectores, no palabras individuales. En cada posición, lm_head genera un vector de logits con 8,192 entradas, una puntuación por cada token del vocabulario, y softmax lo convierte en una distribución de probabilidades sobre todos ellos. El objetivo es un vector one-hot del mismo tamaño, que contiene 1 en la palabra que realmente apareció a continuación y 0 en todos los demás lugares. Por lo tanto, predecir el siguiente token constituye un problema de clasificación cuyas clases son las palabras del vocabulario.
La figura compara los logits en un LLM y en un ViT. Explica un punto sutil: las barras están ordenadas por probabilidad únicamente para facilitar la lectura, mientras que la salida real es un vector con formato fijo, donde la posición i siempre corresponde al identificador del token i. La palabra “son” permanece en el índice 498 en cada paso de procesamiento, ya sea durante el entrenamiento o la inferencia, y softmax solo modifica las alturas de dichas posiciones, nunca su orden. Esto explica por qué el único 1 del objetivo one-hot se encuentra en el índice 498 y no en una posición inicial: el tokenizador asignó a cada token su identidad una sola vez, como su posición correspondiente.
Ese es exactamente el cabezal de clasificación del ViT, que aplica softmax a las puntuaciones de clase y las evalúa mediante entropía cruzada frente a una etiqueta one-hot. Hay dos diferencias: el número de clases (8,192 tokens del vocabulario frente a unas 10 etiquetas de imagen) y la frecuencia (en todas las posiciones T frente a una vez por imagen). Frente a un objetivo one-hot, la entropía cruzada también se reduce a menos el logaritmo de la probabilidad asignada a la palabra correcta, por lo que la pérdida de cada posición se interpreta como una puntuación de sorpresa. “son” con un 9.8 por ciento cuesta 2.32, una palabra que el modelo esperaba firmemente que costara mucho menos, y una verdadera sorpresa costaría mucho más.
Cómo se construye un lote
La instrucción reshape(B*T) indica sencillamente que un lote contiene B secuencias, cada una aportando pérdidas para T posiciones. Entonces, ¿de dónde provienen esas B secuencias y todas tienen la misma longitud?
La gente a menudo pasa por alto un detalle: las secuencias de preentrenamiento no están rellenas y todas tienen exactamente la misma longitud de contexto. El corpus se combina en un único flujo de tokens largo, con un token <|endoftext|> entre los documentos, y se extraen ventanas de longitud fija en posiciones aleatorias. Todas las ventanas provienen de un flujo continuo, por lo que ninguna puede ser corta, y no se necesita ningún token de relleno.
El diagrama muestra el recorrido desde el corpus de entrenamiento hasta los lotes y, finalmente, hasta la pérdida. El token especial en la secuencia es <|endoftext|> (EOT), que funciona como separador en lugar de relleno y indica dónde termina un cuento popular y comienza otro. Una ventana puede cruzar tal límite y contener un EOT a mitad de camino, y así es como el modelo aprende que una historia ha terminado y que comienza una nueva. El relleno sí aparece más tarde, durante el ajuste fino, cuando los ejemplos varían en longitud y las posiciones rellenas se ocultan de la pérdida, lo cual constituye un escenario diferente al del preentrenamiento con datos compactados.
La pérdida es el segundo aspecto en el que el LLM difiere del ViT, y esa diferencia es positiva. El ViT generaba una predicción y una pérdida por imagen. Una secuencia de 128 tokens produce 128 pérdidas, cuyo promedio constituye la pérdida de esa muestra (3.78 para la muestra 1 en la figura); las 32 pérdidas de las muestras en un lote se promedian para obtener la pérdida del lote, de 3.92, que es la que minimiza el algoritmo backpropagation. Gracias a estas dos rondas de promediado, un solo paso reduce el promedio de 32 veces 128, es decir, 4,096 pérdidas por token siguiente, lo que representa una señal mucho mayor por paso que una sola pérdida por imagen.
Observando cómo disminuye la pérdida
Conecte esa pérdida a un optimizador y repita el proceso. Este es el verdadero entrenamiento en el corpus de cuentos.
El gráfico muestra la pérdida de entrenamiento y la pérdida de validación en épocas. El entrenamiento comienza en 9.2, que equivale a ln(8192), la pérdida al adivinar al azar con un vocabulario de 8,192 tokens, y la mejor pérdida de validación alcanza los 4.73. Las curvas se mantienen cercanas, con una diferencia de aproximadamente 1.0, lo que indica que, gracias a un dominio restringido, el modelo generaliza en lugar de simplemente memorizar. El propio modelo es de tamaño reducido.
La figura relaciona el proceso de entrenamiento con la cantidad de pesos aprendibles. Con 9.0 millones de parámetros, el modelo se entrena en unos diez minutos en una GPU de portátil de 6 GB. Un poco más de la mitad se encuentran en los seis bloques transformadores, y el resto pertenece a la tabla de incrustación y a la cabeza de salida. El entrenamiento consiste únicamente en ajustar repetidamente estos 9 millones de valores, de modo que las probabilidades del siguiente token otorgan más importancia al token que realmente apareció a continuación.
En el código, el bucle es el que uno esperaría, con la misma estructura que en un ViT: se toman muestras aleatorias de contexto y ventanas de token siguiente, se ejecuta la pasada forward para obtener una pérdida media, se ponen en cero los gradientes, se realiza la retropropagación y se avanza un paso.
for step in range(train_steps):
xb, yb = get_batch(train_data) # random (context, next-token) windows
_, loss = model(xb, yb) # forward: T predictions -> one mean loss
opt.zero_grad(); loss.backward() # backprop the loss into all 9M parameters
opt.step() # nudge them downhill
El script completo incluye un esquema de tasa de aprendizaje basado en coseno y mecanismos de evaluación.
Qué hay nuevo y qué permanece igual
En comparación con el Vision Transformer, solo hay pequeñas novedades en el LLM: un embedding de tokens que funciona como tabla de consulta, una máscara causal y una pérdida por token siguiente en cada posición que clasifica sobre todo el vocabulario. También destaca una lección práctica: el tokenizador pertenece al modelo, por lo que uno específico para un dominio determinado compensa el esfuerzo invertido. La atención, los embeddings posicionales, la corriente residual, LayerNorm, los bloques MLP y el ciclo completo de retropropagación y optimizador se mantienen sin cambios.
Los números reales revelan el funcionamiento interno: cómo la máscara causal permite que una sola pasada forward calcule las puntuaciones para todas las longitudes de contexto, qué produce attn @ V (cada fila de salida combina los vectores de valor visibles para ese token), cómo se generan los lotes de entrenamiento mediante compactación en lugar de relleno, y cómo la pérdida de una posición, que es una clasificación dentro del vocabulario, se combina con las demás en un promedio único.
Los pesos ya están definidos, pero un modelo entrenado no hace nada hasta que se ejecuta. A continuación viene la inferencia: el bucle de generación, la caché KV que acelera el proceso, y lo que escribe el modelo pequeño al recibir el inicio de una historia.
Lecturas relacionadas
- Reclasificación en RAG: Codificadores transversales vs. evaluadores de LLM y su latencia — Cómo un clasificador de segunda etapa corrige el orden de los candidatos en RAG, comparando codificadores transversales con evaluadores de LLM, costos de latencia, modos de fallo y cómo evaluarlos.
- Flujos de trabajo vs. agentes: ¿Quién debería elegir el siguiente paso de investigación? — Cuándo codificar por defecto un camino de investigación y cuándo dejar que el modelo elija la siguiente búsqueda, con un diseño de herramientas web de LangGraph y Anthropic.