Entrenar un mini-GPT en NumPy

Un mini-GPT recién inicializado reparte su apuesta casi por igual entre las 512 entradas del vocabulario, y la pérdida de la lección sobre el modelo de lenguaje causalBlock 1 · Lesson 1 · In SpanishUna sola columna: el modelo de lenguaje causalEl curso anterior eligió una columna del Transformer, le puso la máscara en todos los bloques y le pidió el token siguiente, posición a posición. Esta lección hace la pregunta que aquél no hizo, qué probabilidad le da ese modelo a un texto entero, y de la respuesta sale todo lo demás: la regla de la cadena convierte esa probabilidad, sin perder nada, en un producto de condicionales; la máscara causal es lo que permite calcular las T de una ventana en una sola pasada; y la pérdida es el logaritmo de ese producto, cambiado de signo y repartido por token. Sin código a propósito: la cuenta se hace a mano sobre una ventana de cuatro palabras, y sale un dieciseisavo por los dos caminos. le cobra por eso $\ln 512 \approx 6.24$ en cada posición. El que cargan las celdas de este bloque paga $3.2$ sobre páginas que nunca vio. Los dos leen los ids del tokenizador de la lección sobre BPEBlock 1 · Lesson 2 · In SpanishBPE de verdad: fusionar, no solo partirEl tokenizador del mini-GPT, escrito desde cero. BPE parte de los 256 bytes con los que se escribe cualquier texto y repite una sola operación: contar los pares de tokens vecinos y fusionar el más frecuente. Esta lección deriva lo que eso garantiza (ningún texto se queda sin tokens, y decodificar devuelve el original byte a byte) y lo que cuesta (cada fusión es una entrada más del vocabulario y ahorra menos tokens que la anterior), y lo comprueba entrenando sobre Marianela las mismas fusiones que usa el modelo. y son la misma red, con $136\,448$ números cada uno. Lo único que los separa es el valor de esos números.

Entre uno y otro hay $2\,000$ pasos de entrenamiento, y esta lección escribe el bucle que los dio. Es el del curso anterior (la pérdida, su gradiente, un paso en contra, repetir) con tres cambios, los mismos que llevan todos los modelos grandes. Cada paso mira unas pocas ventanas sacadas al azar, no el corpus entero. El gradiente no va tal cual al paso: pasa antes por AdamW. Y $\eta$ sube al principio y baja al final. El primero lo derivamos; los otros dos quedan enunciados, con sus razones.

El primero es una cuestión de coste. La parte del corpus con la que se entrenó el checkpoint tiene $125\,789$ tokens, y una ventana de $65$ tokens puede empezar en casi cualquiera de ellos: unas $125\,700$ ventanas solapadas. El gradiente exacto de su pérdida media exige la ida y la vuelta de todas, unos ocho millones de tokens: en el navegador, a algo menos de un milisegundo por token, casi dos horas por paso. Un paso del checkpoint miraba $32$ ventanas de $64$ tokens sacadas al azar, casi cuatro mil veces menos.

Es una encuesta: no le preguntas a todo el país sino a mil personas al azar, y la respuesta sale con ruido pero sin sesgo, con un ruido que baja al crecer la muestra, aunque despacio. La primera sección demuestra las tres cosas.

El segundo cambio lo anunció el curso anterior sin llegar a hacerlo. En su lección sobre el descenso de gradienteDeep Learning for NLP: from the Perceptron to the TransformerBlock 2 · Lesson 6 · In SpanishDescenso de gradienteEl gradiente reúne una derivada parcial por parámetro y apunta, entre todas las direcciones posibles, a la que más sube la pérdida; restarlo un poco cada vez es todo el método. Cuánto es «un poco» decide entre converger, rebotar y divergir, y la cuenta sale entera sobre una parábola., una sola $\eta$ no podía servir a la vez para una dirección empinada y otra casi plana, y los arreglos, momentum y Adam, quedaron fuera. En un batch del mini-GPT hay coordenadas del gradiente miles de veces más pequeñas que otras. AdamW le da a cada parámetro su propia escala, de modo que ninguno se mueve mucho más que $\eta$ por paso, sea cual sea el tamaño de su gradiente.

El tercero, la $\eta$ que sube y baja, se ve en el entrenamiento del checkpoint.

Un batch al azar apunta, de media, adonde apunta el corpus

Fijemos la parte de entrenamiento del corpus y llamemos $\mathcal{L}_i(\theta)$ a la pérdida de la lección sobre el modelo causal en la ventana de $T + 1$ tokens que empieza en la posición $i$: la media de sus $T$ valores de $-\log p_\theta$. La pérdida que entrenar quiere bajar es la media de todas,

donde $i$ se elige al azar, con la misma probabilidad para cada posición en la que cabe una ventana, y $\mathbb{E}$ es la media sobre ese azar.

Un batch de tamaño $B$ son $B$ posiciones $i_1, \dots, i_B$ sacadas así, cada una por su cuenta, y su pérdida es la media de las suyas:

En el código, ventanas(ids, B, T, rng) devuelve X e Y, dos matrices de ids de forma $B \times T$, la segunda corrida una posición, y la red devuelve logits de forma $B \times T \times \lvert V \rvert$. El gradiente de $\hat{\mathcal{L}}$ es el de siempre, el softmax menos el one-hot en cada fila y después la vuelta de la lección sobre backpropagationDeep Learning for NLP: from the Perceptron to the TransformerBlock 2 · Lesson 8 · In SpanishBackpropagation: la derivación completaDerivar peso a peso obliga a recorrer la red una vez por parámetro, y esos recorridos comparten casi todo. Backpropagation calcula en su lugar un solo número por preactivación (el error de esa neurona), con una única pasada del final al principio, y deja que cada peso lea su derivada como el producto de ese error por lo que tenía delante., y perdida devuelve las dos cosas, con los gradientes en un diccionario de las mismas claves y formas que los pesos.

Dos propiedades hacen que $\hat{\mathcal{L}}$ sirva. La primera es que no tiene sesgo. Cada $i_j$ sale del mismo sorteo que $i$, así que la media de $\mathcal{L}_{i_j}(\theta)$ es $\mathcal{L}_{\text{corpus}}(\theta)$ para cada $j$, y la media de $B$ números con esa media tiene esa media. El gradiente es lineal y atraviesa la media igual que atraviesa una suma:

Llamemos $\mathbf{g}$ al lado derecho, el gradiente que costaría dos horas, y $\hat{\mathbf{g}}$ al de dentro de la media, el del batch: dos vectores con una coordenada por parámetro. Un paso con $\hat{\mathbf{g}}$ va, de media, adonde iría un paso con $\mathbf{g}$.

La segunda es que su ruido baja como $1/B$. Lo que $\hat{\mathbf{g}}$ se aleja de $\mathbf{g}$, medido como la media del cuadrado de la distancia, es

el ruido de una sola ventana dividido entre $B$.

El precio está en la raíz. El tamaño típico del error baja como $1/\sqrt{B}$ y el coste del paso crece como $B$: la mitad de ruido cuesta cuatro veces más tokens. Y las ventanas se sacan al azar, no en orden como en el modelo de lenguaje de caracteresDeep Learning for NLP: from the Perceptron to the TransformerBlock 3 · Lesson 7 · In SpanishProyecto: un modelo de lenguaje a nivel de carácterLa RNN vanilla y la BPTT, juntas en una red que predice el siguiente carácter en cada posición: un texto de T caracteres son T ejemplos apilados, y la etiqueta de cada paso es, gratis, el carácter que viene. Se entrena con recorte del gradiente sobre unos miles de caracteres —pocos pasos, para caber en el navegador— y después escribe, realimentando lo que muestrea. La pérdida cae de ln 32 a alrededor de 2.5 y el texto pasa de ruido a español mal hecho —palabras cortas reales, cuasi-palabras largas, ni una frase con sentido—: es lo que aprende un modelo de caracteres, y decirlo sin adornos es la lección., porque las dos cuentas necesitan que cada $i_j$ salga del mismo sorteo que $i$. En orden, el gradiente de cada paso sería el de una página de la novela, y el modelo iría detrás del argumento.

Un paso de AdamW, a la tasa que toca

Numeremos los pasos $s = 1, 2, \dots, S$ y llamemos $\theta_s$ a los parámetros después del paso $s$ (el curso anterior escribía $\theta_t$, pero aquí $t$ ya es la posición dentro de la ventana). El paso $s$ saca un batch, calcula su gradiente $\hat{\mathbf{g}}_s$ y hace con él tres cosas, en este orden.

Primero, el recorte de la lección sobre el gradiente que se desvaneceDeep Learning for NLP: from the Perceptron to the TransformerBlock 3 · Lesson 4 · In SpanishEl gradiente que se desvaneceMultiplicar por la misma matriz recurrente en cada paso vuelve el gradiente exponencial en la distancia: el radio espectral de W_hh decide si se desvanece o explota, la saturación del tanh solo agrava lo primero, y el recorte del gradiente frena la explosión pero no rescata lo desvanecido —que es lo que obligará a cambiar de arquitectura—.: si $\lVert \hat{\mathbf{g}}_s \rVert > 1$, lo dividimos entre su norma. Aquí la norma es la de las $136\,448$ coordenadas juntas y no la de cada grupo de pesos, así que el recorte encoge el paso sin torcer su dirección.

Después, AdamW. Guarda dos medias móviles, una del gradiente y otra de su cuadrado coordenada a coordenada, que empiezan en cero y en cada paso le dan un peso $1 - \rho$ al gradiente nuevo y $\rho$ a lo que ya llevaban:

con $\mathbf{m}_0 = \mathbf{v}_0 = \mathbf{0}$. (El artículo de Adam y minigpt.py llaman $\beta_1$ y $\beta_2$, beta1 y beta2, a lo que aquí es $\rho_1$ y $\rho_2$: la $\beta$ queda reservada para el bloque 2.) Con ellas mueve los parámetros, coordenada a coordenada:

El mini-GPT usa $\rho_1 = 0.9$, $\rho_2 = 0.999$ y $\varepsilon = 10^{-8}$, que sólo evita dividir entre cero. Y $\lambda = 0$: he entrenado el checkpoint sin decaimiento de pesos, así que su Adam es AdamW con la W apagada. Escribirla es el reto de esta lección.

La fórmula queda enunciada (el artículo sobre AdamW de «Para profundizar», al pie, la deriva), pero tres cosas se leen en ella. La división entre $1 - \rho^{s}$ deshace el arranque en cero: en el primer paso, $\mathbf{m}_1 / (1 - \rho_1)$ es $\hat{\mathbf{g}}_1$ y $\mathbf{v}_1 / (1 - \rho_2)$ es su cuadrado, así que el cociente vale el signo del gradiente, coordenada a coordenada. El primer paso mueve $\eta_1$ entero cada parámetro cuyo gradiente no sea cero, sea éste de $10^{-4}$ o de $10$.

Después, el cociente compara la media reciente del gradiente con la raíz de la media reciente de su cuadrado. Si el gradiente insiste en un signo, las dos se parecen y el cociente ronda $\pm 1$; si su signo es ruido, se acerca a $0$. Cada parámetro se mueve, como mucho, del orden de $\eta_s$ por paso: la $\eta$ única del curso anterior, con una escala propia para cada coordenada.

Y $\lambda\,\theta_{s-1}$ va fuera del cociente: cada peso se encoge la misma fracción, $\eta_s \lambda$, tenga la historia de gradientes que tenga. Eso es el decaimiento de pesos desacoplado, la W (weight decay) que Loshchilov y Hutter le añadieron a Adam. Sumado a $\hat{\mathbf{g}}_s$, pasaría por la división y encogería menos a los pesos que más se mueven.

Queda $\eta_s$. El checkpoint la saca de un calendario de dos tramos, un calentamiento (warm-up) lineal durante $S_{\text{cal}}$ pasos y un descenso en coseno hasta la décima parte,

con $\eta_{\max} = 3 \cdot 10^{-3}$, $S_{\text{cal}} = 200$ y $S = 2\,000$: la curva de arriba de la figura.

El calentamiento sale de la primera lectura de la fórmula. Mientras $\mathbf{m}_s$ y $\mathbf{v}_s$ llevan uno o dos gradientes, cada parámetro da un paso de casi $\eta_s$ en la dirección del signo de un gradiente lleno de ruido: $136\,448$ pasos a ciegas del mismo tamaño, que sólo son pequeños si $\eta_s$ lo es.

El descenso sale de la primera sección. Cerca del final, con la pérdida casi plana, $\mathbf{g}$ es pequeño pero el ruido del batch no se va, y el gradiente de cada paso es casi todo ruido. Con $\eta$ fija, el modelo tiembla alrededor del mínimo a una distancia que crece con $\eta$; bajarla al final lo deja asentarse.

El mini-GPT, entrenado delante de ti

Las tres celdas corren sobre el mini-GPT: minigpt.py, el fichero de NumPy que produjo el checkpoint, y minigpt.json, sus pesos. Es la columna del proyecto del TransformerDeep Learning for NLP: from the Perceptron to the TransformerBlock 5 · Lesson 9 · In SpanishProyecto: un Transformer desde ceroEl artículo entero, montado y ejecutado: un vocabulario de doce entradas compartido por los dos idiomas, dos bloques por columna, dos cabezas y una frase de cinco tokens que sale por arriba convertida en seis distribuciones. Las seis funciones que hacen el trabajo son los desafíos de las lecciones anteriores sin tocar una letra, y lo que las une es que todo lo de en medio recibe T por d_model y devuelve T por d_model. Nadie ha entrenado estos pesos, así que la traducción que propone no significa nada; lo que sí se puede comprobar con pesos cualesquiera son tres cosas, y la segunda es la que este montaje añade: que ninguna posición del decoder ve lo que todavía no se ha escrito exige la máscara en cada uno de los bloques, no en el primero, y quitar una sola se mide. en pequeño (dos bloques, cuatro cabezas, $d_{\text{model}} = 64$) más perdida, que hace la ida, la pérdida y la vuelta, y Adam.

La primera carga el tokenizador, el modelo y el corpus, aparta el último 10 %, el texto reservado que el checkpoint no vio nunca, y saca un batch de cuatro ventanas de ocho tokens para que veas las formas. Después mide sobre él el checkpoint y la misma red con los pesos del paso 0.

La primera fila de X es Anduvieron b en ocho tokens, y la de Y, la misma corrida uno, acaba en re. La red sin entrenar paga $6.235$, casi $\ln 512 = 6.238$: con pesos de tamaño $0.02$ los logits salen casi iguales y el softmax reparte por igual. El checkpoint paga $3.180$. Entre los dos números están los $2\,000$ pasos de la figura.

La segunda pone a prueba el $1/B$. Para $B = 1$, $4$ y $16$ ventanas de ocho tokens, saca ocho batches independientes y mide cuánto se alejan sus gradientes de la media de los ocho: una estimación de $\mathbb{E}\left[\lVert \hat{\mathbf{g}} - \mathbf{g} \rVert^{2}\right]$. Si la sección tiene razón, $B$ por el ruido no cambia.

El ruido cae de $540$ a $155$ y a $38$, catorce veces menos con dieciséis veces más ventanas, y $B$ por el ruido se queda entre $540$ y $620$. Las estimaciones tienen su propio ruido; cambia la semilla y la constante sigue ahí.

La tercera sigue entrenando, sobre una copia del checkpoint y con un Adam nuevo: el checkpoint guarda $\theta$ pero no $\mathbf{m}$ ni $\mathbf{v}$, así que el calentamiento vuelve a hacer falta. Entrena sobre los últimos $3\,000$ tokens del texto reservado, las últimas páginas de Marianela, con una sola ventana de $16$ tokens por paso: el batch más ruidoso, pero el que cabe en el navegador. El calendario es el del checkpoint en pequeño, con $\eta_{\max} = 5 \cdot 10^{-4}$, más o menos donde el suyo lo dejó ($4 \cdot 10^{-4}$ en el paso 1 750). En un dispositivo lento, la celda se detiene antes de que el navegador la corte, y dice dónde.

La pérdida sobre esas páginas baja de $3.473$ a $3.104$, pero no en cada medida: sube en el paso 50, y otra vez en el 100 y el 125. Con una ventana por paso, cada $\hat{\mathbf{g}}_s$ es casi todo ruido y el modelo avanza dando tumbos. Y lo que baja es la pérdida sobre las páginas con las que entrena: si eso es aprender español o aprenderse esas páginas lo mide la lección sobre perplejidad, más adelante en el bloque.

Ahora quítale el calentamiento: pon eta_max, S_cal, S = 3e-3, 1, 200, la tasa máxima del checkpoint sin calentar, y vuelve a ejecutarla. En 25 pasos la pérdida salta de $3.47$ a $4.89$, porque el primer paso mueve cada parámetro $3 \cdot 10^{-3}$ en la dirección de un solo gradiente lleno de ruido. Con S_cal = 50 el salto se retrasa, no desaparece: llega hacia el paso 75, a $4.35$, justo después de que la tasa toque su máximo. El calentamiento protege el arranque y nada más. El checkpoint salió del final de su calendario, con la tasa en la décima parte, y reanudarlo a la máxima lo saca de donde estaba.

Comprueba tu intuición

Cuatro preguntas sobre el batch y el primer paso de Adam, y un reto: escribir AdamW.

X, Y = ventanas(ids, 4, 16, np.random.default_rng(0))
print(X.shape, Y.shape, (X[:, 1:] == Y[:, :-1]).all())
 

El mini-GPT ya se entrena, y lo que sale de él es lo mismo que salía en la tabla a mano de la lección sobre el modelo de lenguaje causalBlock 1 · Lesson 1 · In SpanishUna sola columna: el modelo de lenguaje causalEl curso anterior eligió una columna del Transformer, le puso la máscara en todos los bloques y le pidió el token siguiente, posición a posición. Esta lección hace la pregunta que aquél no hizo, qué probabilidad le da ese modelo a un texto entero, y de la respuesta sale todo lo demás: la regla de la cadena convierte esa probabilidad, sin perder nada, en un producto de condicionales; la máscara causal es lo que permite calcular las T de una ventana en una sola pasada; y la pérdida es el logaritmo de ese producto, cambiado de signo y repartido por token. Sin código a propósito: la cuenta se hace a mano sobre una ventana de cuatro palabras, y sale un dieciseisavo por los dos caminos.: en cada posición, una distribución sobre las 512 entradas. La pérdida sólo le pregunta cuánto apostó por el token que de verdad venía; nunca le pide que elija uno. Para escribir texto hay que elegir, token a token, y el modelo no dice cómo.

Ése es el asunto de la lección siguiente, sobre el muestreo: quedarse siempre con la entrada más probable, que en el mini-GPT acaba repitiendo no, no, no, o sortear entre las posibles con una temperatura, un top-k o un top-p, y qué parte de la distribución tira cada uno.