Saltar al contenido
← Attention Is All You Need

Lección 7 de 10 · La arquitectura completa · 22 min

Entrenamiento, resultados y legado

La receta de entrenamiento del paper (warmup, label smoothing, BPE), los resultados en WMT 2014 y cómo el Transformer dio origen a BERT, GPT, ViT y a la carrera contra el costo cuadrático.

Notebook del episodio

El código de esta lección, para correrlo y modificarlo en Google Colab, gratis y con tu cuenta. Para guardar tus cambios: Archivo → Guardar una copia en Drive.


Datos y tokenización

El paper entrena en las tareas de traducción de WMT 2014:

  • Inglés → alemán: unos 4,5 millones de pares de oraciones, tokenizados con byte-pair encoding (BPE) y un vocabulario compartido fuente-objetivo de unos 37.000 tokens.
  • Inglés → francés: unos 36 millones de pares, con un vocabulario de 32.000 word-pieces.

BPE (Sennrich et al., 2016) empieza con caracteres y fusiona iterativamente los pares más frecuentes hasta alcanzar el tamaño de vocabulario deseado. Las palabras comunes quedan como un solo token y las raras se descomponen en sub-palabras ("des", "atención", "ado"). Así el modelo nunca enfrenta una palabra completamente desconocida, y un vocabulario compartido permite que nombres propios y cognados se copien entre idiomas. Hoy alguna variante de tokenización por sub-palabras es universal en los LLM.

La receta de optimización

El entrenamiento usa Adam con β1=0,9\beta_1 = 0{,}9, β2=0,98\beta_2 = 0{,}98 y ϵ=10−9\epsilon = 10^{-9}, y un learning rate que cambia con cada paso:

lrate=dmodel−0.5⋅min⁡ ⁣(step−0.5,  step⋅warmup_steps−1.5)lrate = d_{\text{model}}^{-0.5} \cdot \min\!\left(step^{-0.5},\; step \cdot warmup\_steps^{-1.5}\right)

con warmup_steps=4000warmup\_steps = 4000. La fórmula tiene dos regímenes:

  • Warmup: durante los primeros 4000 pasos, el segundo término es el menor y el learning rate crece linealmente.
  • Decaimiento: después, domina el primero y el learning rate decae como 1/step1/\sqrt{step}.

Ambos términos se igualan justo en step=warmup_stepsstep = warmup\_steps, que es donde se alcanza el máximo. Con dmodel=512d_{\text{model}} = 512, ese pico es aproximadamente 512−0.5⋅4000−0.5≈7×10−4512^{-0.5} \cdot 4000^{-0.5} \approx 7 \times 10^{-4}.

python
def transformer_lr(step: int, d_model: int = 512, warmup: int = 4000) -> float:
    step = max(step, 1)
    return d_model ** -0.5 * min(step ** -0.5, step * warmup ** -1.5)

# Pico en step = 4000: ≈ 6.99e-4

Comprueba tu comprensión

En el schedule del paper con warmup_steps = 4000, ¿qué pasa con el learning rate después del paso 4000?

Regularización

El paper usa dos técnicas de regularización:

Dropout residual. Se aplica dropout a la salida de cada sub-capa antes de sumarla al residual, y también a la suma de embeddings y codificaciones posicionales. Para el modelo base, Pdrop=0,1P_{drop} = 0{,}1.

Label smoothing. En vez de entrenar contra la distribución one-hot del token correcto, se reparte una masa ϵls=0,1\epsilon_{ls} = 0{,}1 de manera uniforme sobre el resto del vocabulario. El objetivo deja de ser "probabilidad 1 a la respuesta correcta" y pasa a ser "probabilidad alta, pero no absoluta".

“This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score.”
— Sección 5.4

Es un detalle revelador: la métrica que optimizas (perplejidad) y la que te importa (calidad de traducción) no siempre van en la misma dirección. Un modelo menos sobreconfiado generaliza mejor y produce mejores búsquedas en beam search.

Hardware y costo

Todo se entrenó en una máquina con 8 GPU NVIDIA P100:

  • Modelo base (6 capas, dmodel=512d_{\text{model}} = 512, 8 cabezas): 100.000 pasos, unas 12 horas.
  • Modelo grande (dmodel=1024d_{\text{model}} = 1024, dff=4096d_{ff} = 4096, 16 cabezas): 300.000 pasos, unos 3,5 días.

Para la inferencia usaron beam search con tamaño de haz 4 y penalización por longitud α=0,6\alpha = 0{,}6, y promediaron los pesos de los últimos checkpoints.

El paper estima el costo de entrenamiento en FLOPs: del orden de 3,3×10183{,}3 \times 10^{18} para el modelo base y 2,3×10192{,}3 \times 10^{19} para el grande, una fracción del costo de los mejores modelos anteriores. Para dimensionar: los LLM de frontera actuales se entrenan con varios órdenes de magnitud más cómputo.

Resultados

En WMT 2014 inglés → alemán, el Transformer grande logró 28,4 BLEU, superando por más de 2 puntos a los mejores resultados previos, incluyendo ensambles. Incluso el modelo base superó a todos los modelos individuales y ensambles publicados, a una fracción del costo de entrenamiento.

En inglés → francés, el modelo grande alcanzó 41,8 BLEU, un nuevo estado del arte para un modelo individual, también con un costo de entrenamiento mucho menor que el de los competidores.

Para mostrar que la arquitectura no estaba sobreajustada a la traducción, los autores también la aplicaron a análisis sintáctico de constituyentes (constituency parsing) en inglés, con resultados competitivos y poca adaptación específica.

Comprueba tu comprensión

¿Qué resultado obtuvo el Transformer grande en WMT 2014 inglés → alemán?

Los estudios de ablación

La Tabla 3 del paper es una mina de información. Algunas conclusiones:

  • Número de cabezas: una sola cabeza es claramente peor; demasiadas cabezas también degradan. Hay un punto medio.
  • Reducir dkd_k perjudica la calidad, lo que sugiere que calcular compatibilidad no es trivial.
  • Modelos más grandes son mejores, y el dropout es clave para evitar sobreajuste.
  • Posiciones aprendidas vs. sinusoidales: resultados casi idénticos.

La observación "más grande es mejor" pasa casi desapercibida en el paper, pero anticipa la década siguiente: los Transformers escalan de manera notablemente predecible, algo que los trabajos sobre leyes de escalamiento (scaling laws) formalizarían años después.

El legado

Pocas arquitecturas tuvieron un impacto comparable. En pocos años, el Transformer dejó de ser "un modelo de traducción" y se convirtió en la arquitectura por defecto del deep learning:

  • BERT (2018): solo el encoder, preentrenado con masked language modeling sobre texto masivo. Redefinió el estado del arte en comprensión de lenguaje e instauró el paradigma preentrenamiento + ajuste fino (fine-tuning).
  • GPT (2018) → GPT-2 → GPT-3 y sucesores: solo el decoder, preentrenado para predecir el siguiente token. Al escalar, aparecieron capacidades como el aprendizaje en contexto (in-context learning). Es el linaje de los LLM conversacionales actuales.
  • T5 (2019): mantuvo el encoder-decoder y formuló todas las tareas como texto → texto.
  • Vision Transformer, ViT (2020): cortó imágenes en parches de 16×16, trató cada parche como un token y demostró que, con suficientes datos, la atención compite o supera a las CNN en visión.
  • Y más allá: audio (Whisper), proteínas (el módulo de atención de AlphaFold 2), modelos multimodales, robótica.

La deuda pendiente: O(n2)O(n^2)

El precio de conectar todo con todo es que la atención cuesta O(n2)O(n^2) en tiempo y, en una implementación ingenua, en memoria. Con oraciones de 50 tokens no importa; con documentos de 100.000 tokens es el factor dominante. Eso dio lugar a varias líneas de trabajo:

  • Atención dispersa (sparse): Sparse Transformer, Longformer y BigBird atienden solo a un subconjunto estructurado de posiciones (ventanas locales, tokens globales).
  • Aproximaciones lineales: Linformer, Performer y la atención lineal aproximan el softmax para lograr costo O(n)O(n), con cierta pérdida de calidad.
  • Atención exacta pero eficiente en IO: FlashAttention (2022) no cambia la matemática, sino cómo se mueve la memoria en la GPU: calcula la atención por bloques sin materializar la matriz n×nn \times n completa. Sigue siendo O(n2)O(n^2) en FLOPs, pero lineal en memoria y mucho más rápida. Hoy es estándar.
  • Inferencia: el KV cache, multi-query y grouped-query attention reducen la memoria de las keys y values durante la generación.
  • Alternativas a la atención: modelos de espacio de estados (state space models) como Mamba, y arquitecturas híbridas que combinan capas recurrentes eficientes con algunas capas de atención.

Es una ironía histórica: el paper abandonó la recurrencia por su secuencialidad, y algunas de las alternativas más prometedoras para contextos largos recuperan ideas recurrentes, pero diseñadas para entrenarse en paralelo.

Comprueba tu comprensión

¿Qué hace FlashAttention para acelerar la atención?

Cierre del curso

Recorriste el paper completo: el cuello de botella que lo motivó, la atención escalada y su argumento de varianza, las cabezas múltiples, las rotaciones que codifican posición, la arquitectura con sus máscaras y la receta que la hizo entrenable. Con esto tienes la base para leer prácticamente cualquier paper de LLM moderno: casi todos son variaciones sobre las piezas que acabas de construir.

Un buen ejercicio final: toma el código de las lecciones, ensámblalo en un Transformer decoder-only chico y entrénalo para predecir el siguiente carácter de un texto en español. Ver aparecer palabras con sentido a partir de ruido es la mejor forma de convencerte de que lo entendiste.

¿Qué tal esta lección?

Sé el primero en calificarla

Entra para calificar

¿Terminaste la lección?

Márcala para sumar +50 XP y mantener tu racha.