La tasa de aprendizaje es el hiperparámetro que rompe producción

La parte 8 terminó con un gradiente, que indica hacia dónde mover cada peso. Cuánto moverlo lo decide el optimizador, y esa decisión se reduce a un puñado de números en un fichero de configuración que nada en todo el stack comprueba.

He leído el paso de Adam en torch y el schedule de nanoGPT, y después he entrenado de 21 maneras el modelo de caracteres de la parte 8 (cuatro bloques que leen el texto carácter a carácter): siete tasas de aprendizaje, cada una sin warmup, con warmup y con warmup más un decaimiento coseno. Ninguna de las 21 ejecuciones se cayó.

Tres optimizadores en cinco líneas

El descenso de gradiente a secas resta el gradiente multiplicado por un tamaño de paso. El momentum va acumulando los gradientes pasados en una suma que se encoge un poco en cada paso, y es esa suma la que mueve el peso. Así, un peso empujado dos veces hacia el mismo lado coge velocidad. Adam añade una segunda media, la del gradiente al cuadrado, y divide por su raíz cuadrada. El paso de torch 2.14, condensado a partir de adam.py:456-546, queda así.

exp_avg.lerp_(grad, 1 - beta1)                    # average of the gradient
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)   # average of its square
step_size = lr / (1 - beta1 ** step)              # bias correction, early steps
denom = exp_avg_sq.sqrt() / (1 - beta2 ** step) ** 0.5 + eps
param.addcdiv_(exp_avg, denom, value=-step_size)  # one step per weight

Toda la idea está en esa división. A un peso cuyo gradiente ha sido grande últimamente le toca un divisor grande, y a uno con gradientes pequeños, uno pequeño. Cada peso acaba así con su propio tamaño de paso, marcado por su historia reciente, y lr solo pone la escala. Otros dos números de la configuración, beta1 y beta2, fijan lo rápido que cada media olvida lo que ha visto (nanoGPT usa 0,9 y 0,95).

AdamW añade una línea, param.mul_(1 - lr * weight_decay) en adam.py:419, que encoge un poco cada peso en cada paso. La última sección trata de dónde se coloca esa línea.

La escala de la pérdida se cancela y lo paga la memoria

En la parte 8, multiplicar una pérdida por 40 no cambió nada, y ahora se ve por qué. Si multiplicas cada gradiente por 40, exp_avg se multiplica por 40 y exp_avg_sq por 1600, así que su raíz cuadrada también se multiplica por 40 y el factor se cancela en la división. Se cancela en todas partes salvo en eps, la constante que se suma al denominador para que no se acerque a cero. Escalar la pérdida por 40 equivale a dividir esa constante entre 40. Si alguna vez has cambiado las unidades de una métrica y las alertas han seguido igual porque todos los umbrales eran relativos, aquí pasa lo mismo.

En tres pasos del modelo de caracteres, el desplazamiento mediano de los pesos con AdamW fue de 0,0022 tanto con la pérdida escalada como sin escalar, y los dos vectores de actualización apuntaban en la misma dirección. Con SGD a secas, el mismo cambio multiplicó por 17,5 el paso mediano y desvió 30 grados la actualización. Eso sí, la cancelación no es idéntica peso a peso: la mayor discrepancia fue de 0,00134, más de la mitad de un paso típico.

Lo que Adam gana en robustez lo paga en memoria. Guarda dos números por peso, que torch llama exp_avg y exp_avg_sq, mientras que SGD con momentum guarda uno y SGD a secas, ninguno. Lo he comprobado en el estado del optimizador de torch con el modelo de caracteres, y salen justo esos números. Para GPT-2 small en fp32 son 996 MB de estado del optimizador, junto a 498 MB de pesos y otros 498 MB de gradientes. Pongamos 2 GB antes de contar una sola activación del grafo de la parte 8.

Con AdamW, las tres barras que no se mueven con el batch suman los 2 GB de antes. Cambia a SGD y verás desaparecer un tercio del total. Luego sube el batch y mira qué le hacen las activaciones a ese mismo presupuesto.

El primer paso mide exactamente la tasa de aprendizaje

Vuelve a mirar step_size = lr / (1 - beta1 ** step). En el primer paso, exp_avg es el gradiente multiplicado por 1 - beta1, justo lo que divide la corrección. Al cancelarse los dos, la primera actualización se queda en el gradiente dividido por su propio tamaño, multiplicado por lr.

Con lr a 0,001, el desplazamiento mediano de los pesos en el primer paso fue de 0,001, exactamente la tasa de aprendizaje, fuera cual fuera el gradiente. El tamaño del gradiente no llega nunca a la actualización, y de él solo sobrevive el signo, peso a peso. A partir del segundo paso, lo que sobrevive es cada gradiente comparado con su propio tamaño reciente.

Para que salga exacto hacen falta dos condiciones, y las dos dependen de la configuración. Un gradiente cercano a eps (10⁻⁸ por defecto) lo estropea: con 10⁻⁹ el primer paso se queda en una décima parte de la tasa, y con 10⁻¹², en una diezmilésima. Además, el AdamW de torch aplica por defecto un weight decay de 0,01 si no le pasas otro valor. De ahí que el primer paso más grande que he medido fuera de 0,00101, apenas por encima de la tasa.

Para eso está el warmup. Si la tasa arranca cerca de cero y sube de forma lineal durante los primeros cientos de pasos, esas primeras actualizaciones, que miden exactamente la tasa, salen pequeñas. El schedule de nanoGPT, en train.py:231-242, cabe en ocho líneas.

if it < warmup_iters:                       # 1) linear warmup
    return learning_rate * (it + 1) / (warmup_iters + 1)
if it > lr_decay_iters:                     # 2) then a floor
    return min_lr
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
assert 0 <= decay_ratio <= 1
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))   # 3) cosine in between
return min_lr + coeff * (learning_rate - min_lr)

Con 100 pasos de warmup, la primera actualización se quedó en 0,0000099 frente a 0,001 (cien veces menor) y a partir de ahí fue creciendo. Pasado el warmup, el resto del schedule va bajando la tasa hasta una décima parte del pico, que es lo que fija el min_lr de nanoGPT en train.py:68. El único assert del schedule comprueba la aritmética del cociente, y el valor que has configurado ni lo mira.

Qué dice el barrido

He entrenado el modelo de caracteres durante 600 pasos con tres schedules (ninguno, 100 pasos de warmup y warmup más el decaimiento coseno) y siete tasas de aprendizaje para cada uno, de 0,0001 a 0,1. Los gradientes iban recortados a 1,0 en todo momento, como hace nanoGPT. El barrido no puede separar el warmup del recorte, y el recorte ya estaba haciendo su parte. Con 0,003 y sin warmup, quitar el recorte llevó la puntuación de 2,18 a 2,38 y dejó que la norma del gradiente llegara a 79. Como referencia, adivinar solo por frecuencia da 3,23 en el texto de test.

Ninguna ejecución dio un NaN, ni siquiera con una tasa de aprendizaje diez veces mayor que la mejor.

La mejor ejecución sacó 1,73 con warmup, coseno y una tasa de 0,01. La misma tasa dio 1,84 solo con warmup y 2,61 sin ninguno de los dos. Sin warmup, ninguna tasa bajó del 2,02 que sacó la de 0,001. El warmup, por tanto, no rescató ninguna ejecución rota. Hizo viable una tasa diez veces mayor, y en esa tasa estaba la buena puntuación.

Con 0,1 constante, la pérdida se estancó cerca de 3,0, apenas mejor que adivinar por frecuencia de letras. Bajar, bajó, pero luego dejó de bajar. Ese valor solo te parecería malo si supieras lo que hace el mismo modelo con 0,01. Con el coseno, que baja la tasa hasta 0,01, la ejecución de 0,1 siguió mejorando hasta 2,71, aunque se quedó a casi un nat de la mejor.

Empieza en 0,01 y pulsa sin warmup, que también quita el coseno, y verás que la puntuación en test sube de 1,73 a 2,61. Luego arrastra hasta 0,1 y mira cómo la curva se aplana sin llegar a explotar, con el indicador de NaN todavía en «nunca».

Para ser justos con lo que da de sí un portátil, un modelo de cuatro bloques entrenado en él no es un preentrenamiento. Como cada cifra de aquí sale de una sola semilla, por debajo de unos 0,05 nats no hay nada que leer. A escala real, una tasa mala sí produce el pico de pérdida y el NaN. Lo que se transfiere es la forma: la misma tasa que va bien con warmup desperdicia el entrenamiento sin él, y el fallo que tienes que saber detectar no es que la ejecución se caiga.

El weight decay y dónde ponerlo

El weight decay encoge los pesos hacia cero. La forma original consiste en sumar weight_decay * param al gradiente, con lo que ese término pasa por la división de Adam y se reescala peso a peso como todo lo demás. El param.mul_(1 - lr * weight_decay) de AdamW lo saca de ese camino, y de ahí viene la W.

La diferencia salta a la vista. Con el mismo entrenamiento de 400 pasos y un weight decay de 0,1 en todos los parámetros, la versión desacoplada sacó 2,22 en el texto de test y la acoplada 3,03, apenas mejor que adivinar por frecuencia. La acoplada terminó con una norma de los pesos (la longitud de todos ellos tomados como un solo vector) de 21, frente a 41. Reescalado por Adam, el weight decay fue lo bastante fuerte para vaciar el modelo.

Hay dos cosas que decir de esa cifra. La primera es que la misma ejecución sin weight decay sacó 2,21. En 400 pasos, el weight decay desacoplado no aportó nada y toda la diferencia es el daño que hace el otro camino. La segunda es que aplicar weight decay a todos los parámetros probablemente empeora el caso acoplado, porque la parte del gradiente que encoge los pesos cae también sobre las escalas de normalización, que empiezan en 1,0 y en la inicialización concentran la mayor parte de la norma de los pesos.

nanoGPT aplica el weight decay de forma selectiva en model.py:269-274, donde lo llevan las matrices pero no los sesgos ni las escalas de LayerNorm. Son dos líneas de filtrado, y saltárselas no da ningún error.

Qué vigila el ingeniero de ML

Son cuatro hábitos, y los dos primeros consisten solo en guardar números.

Registra la tasa de aprendizaje efectiva: con la configurada no basta. El número que importa es el que ha devuelto el schedule en este paso. Un schedule que se reinició al reanudar repitió el warmup, arrancó 70 veces por debajo de la tasa de una reanudación completa y acabó la ejecución más de tres veces por encima, mientras el fichero de configuración decía lo mismo en las dos ejecuciones. En los mismos 150 pasos sacó 2,275 frente a 2,230, dentro del ruido de una sola semilla. Solo la tasa registrada lo habría delatado.

Guarda en el checkpoint el estado del optimizador, además de los pesos. En GPT-2 small, las dos medias de AdamW ocupan 996 MB, y si las tiras, el promediado vuelve a empezar sin avisar. Sobre el mismo batch, el primer paso tras reanudar dio 2,64 frente a 2,38 con una reanudación completa. Mi ejecución se recuperó en 150 pasos hasta quedar a 0,007 de la reanudación completa, de modo que a esta escala el estado no valía casi nada. La razón para guardarlo está en el transitorio, y no he probado una tasa lo bastante grande como para que resultara fatal.

Dimensiona la máquina con 16 bytes por parámetro, sin contar las activaciones. En fp32, pesos, gradientes y las dos medias de AdamW ocupan cuatro bytes por parámetro cada uno, así que si dimensionas solo por los pesos te quedas corto en tres cuartas partes. La precisión mixta no rebaja la suma: pesos y gradientes de 2 bytes, más una copia maestra de 4 bytes y las dos medias, siguen sumando 16 bytes. En ese caso, si dimensionas solo por los pesos, te quedas corto en siete octavos.

Haz el barrido de la tasa en un modelo pequeño y fija una regla para abortar. El modelo pequeño te da la forma, pero el valor no se transfiere, porque la mejor tasa baja a medida que crece el modelo, salvo que uses algo como μP. Como nada lanza una excepción, la regla tiene que ponerla tú. La mía saltaba cuando la pérdida llegaba al doble de su mejor valor, y lo hizo en cuatro de las 21 ejecuciones. Tres eran de 0,1; la cuarta siguió adelante y terminó bien. Una regla tan barata matará de vez en cuando una ejecución buena, y ese es el precio de pillar las malas en los diez primeros pasos.

Esto tampoco es trabajo de modelado. Son cinco líneas de aritmética, un schedule de ocho líneas y un fichero de configuración que nada valida. Entre mi mejor ejecución y la peor hubo 1,3 nats con el mismo modelo, los mismos datos y la misma semilla, y lo único que cambiaba eran tres números de la configuración.