Pérdida y gradientes sin fórmulas: la retropropagación es recorrer un grafo

La parte 7 acabó con cuatro cambios en un bloque, cada uno grabado en los pesos, y todos esos pesos se aprendieron. Esta parte cuenta cómo: hace falta un único número que diga cuánto se equivoca el modelo y, para cada peso, saber hacia dónde moverlo para que ese número baje. Esa dirección, peso a peso, es el gradiente.

He leído el engine.py de micrograd y lo he contrastado con torch. micrograd, que Andrej Karpathy escribió para enseñar diferenciación automática, hace todo el trabajo con números sueltos en 94 líneas. Después he entrenado un modelo pequeño de cinco maneras, cuatro de ellas mal, y ninguna de las malas ha dado error.

La pérdida es la sorpresa media

Un modelo de lenguaje devuelve una probabilidad para cada token del vocabulario. La pérdida (la entropía cruzada) solo se fija en el token que venía a continuación en el texto y toma el logaritmo de su probabilidad cambiado de signo. Si el modelo le da probabilidad 1, cuesta 0; si le da un 1 %, cuesta 4,6. A ese coste lo llamo sorpresa, y se mide en nats porque el logaritmo es en base e, que vale unos 2,718.

He pasado 1024 tokens de una entrada anterior por GPT-2 small y la sorpresa media ha salido 3,595 (el 3,60 que daba la parte 5). Si elevas e a esa potencia sale 36, es decir, de media el modelo dudaba tanto como si eligiera al azar entre 36 tokens. Una predicción uniforme entre 50 257 tokens da 10,8.

La media esconde una distribución asimétrica, con una forma que conoce bien quien haya vigilado alguna vez una latencia p99. El token mediano costó 2,85 y la peor décima parte de los tokens se llevó el 28 % de la pérdida. Los más sorprendentes eran nombres: " Claude" después de «renaming the assistants alone drops» costó 16,0. Entre los que menos sorprendieron estaban los trozos de una URL, como «org» después de «arxiv.», que costó una diezmilésima de nat.

Al entrenar, esa sorpresa se promedia sobre todas las posiciones de todas las secuencias de un batch hasta quedar en un solo número. Todo lo que sigue trata de cómo volver desde ese número hasta los pesos.

La pasada hacia delante construye un grafo

micrograd envuelve cada número en un Value que recuerda de dónde viene. Así es la multiplicación, condensada de engine.py:24-33.

def __mul__(self, other):
    out = Value(self.data * other.data, (self, other), '*')
    def _backward():
        self.grad += other.data * out.grad    # d(out)/d(self) is other
        other.grad += self.data * out.grad    # d(out)/d(other) is self
    out._backward = _backward
    return out

Cada operación hace dos cosas: calcula su resultado y deja tras de sí un nodo que conoce sus entradas y un closure que conoce su propia derivada local. La pasada hacia delante calcula una respuesta y, como efecto secundario, construye un grafo. Si alguna vez has escrito un tracer que registra cada llamada enlazada a las llamadas que le dieron sus datos, ya tienes hecha la mitad.

Una expresión pequeña, relu(a·b + c)·a, construye 7 nodos, y una red de 41 parámetros evaluada con cuatro ejemplos, 394. torch hace lo mismo con tensores. Cada nodo guarda lo que vaya a necesitar en la pasada hacia atrás. En GPT-2 small, una secuencia de 1024 tokens deja unos 820 MB de tensores guardados, más que los 498 MB que ocupan los pesos del modelo. Con un batch de ocho son 6,5 GB.

La pasada hacia atrás desanda el grafo

Value.backward, condensado de engine.py:54-70, es todo el algoritmo.

topo, visited = [], set()
def build_topo(v):                     # depth-first: children before parents
    if v not in visited:
        visited.add(v)
        for child in v._prev:
            build_topo(child)
        topo.append(v)
build_topo(self)
self.grad = 1
for v in reversed(topo):               # the output first, the inputs last
    v._backward()

Hace una ordenación topológica y luego recorre la lista al revés. Ese orden garantiza que cada nodo haya reunido todas sus contribuciones antes de transmitir nada. El closure de cada nodo coge el gradiente que le ha llegado, lo multiplica por su derivada local y suma el resultado a sus entradas. Eso es la regla de la cadena hecha bucle, y no hay nada más. Si has escrito un sistema de build que recorre un grafo de dependencias, esto es el mismo recorrido hacia atrás: ningún paso se ejecuta hasta que todos los que dependen de él han entregado su parte.

Recorre la pasada hacia atrás nodo a nodo y después cambia a grad = para ver cómo a pierde una de sus dos contribuciones.

Todo depende de ese +=. En la expresión de arriba a se usa dos veces, así que le llega gradiente desde dos nodos. Con += recibe −2, que es lo que da torch; con =, la segunda escritura machaca la primera y a se queda con −6. No salta ningún error y el número, sencillamente, está mal. En la red de 41 parámetros, los gradientes de micrograd y los de torch se diferenciaron como mucho en 0,000007.

torch acumula igual, en el .grad de cada peso, hasta que algo lo pone a cero. nanoGPT, el GPT-2 mínimo que esta serie viene leyendo, usa esa acumulación a propósito. En train.py:292-305 procesa 40 batches pequeños en una GPU, divide cada pérdida entre 40 y llama 40 veces a backward sobre los mismos buffers antes de hacer una sola actualización. Luego optimizer.zero_grad, en train.py:314, los vacía.

Con cadenas largas, el gradiente se desvanece o explota

En cualquier camino de la pérdida a un peso, la regla de la cadena multiplica entre sí las derivadas locales, una por nodo. Las redes profundas tienen caminos largos. El producto de muchos números un poco por debajo de 1 es diminuto, y el de muchos un poco por encima, enorme. Eso es, en pura aritmética, el desvanecimiento y la explosión del gradiente, y para verlo no hace falta saber cálculo.

He montado pilas de capas tanh de 256 de ancho y sin conexión residual; tanh aplasta cualquier valor al intervalo entre −1 y 1. Lo único que cambiaba de una pila a otra era cuánto escalaba cada capa su entrada. Con una escala de 0,5 y 30 capas, el gradiente que llegaba a la primera era 7 × 10⁻⁹ veces el de la última, y con 60 capas, unas 10⁻¹⁷ veces. Con una escala de 2 y 60 capas, la primera recibía unas 7700 veces más.

El flujo residual de la parte 5, el bus del que lee cada bloque y al que suma su resultado, resuelve la mitad del problema, la del desvanecimiento. Si cada capa es residual, x + layer(x), todas tienen un camino directo de vuelta a la pérdida. En las mismas pilas, la primera capa no recibió nunca menos gradiente que la última, aunque con 60 capas y escala 2 seguía recibiendo 992 veces más. GPT-2 small está construido así, y al medirlo con una secuencia ningún bloque recibió un gradiente raquítico ni desbocado.

Elige 60 capas, pon la escala en 0,5 y verás cómo se desvanece el gradiente de la pila simple. Si después pasas a la residual, con 0,5 no se desvanece nada y con 2 la primera capa sigue recibiendo unas mil veces más.

Para eso está el recorte del gradiente. clip_grad_norm_ mide la longitud del gradiente entero, con todos los pesos a la vez, y si pasa de un límite lo reduce todo en la misma proporción. nanoGPT pone el límite en 1,0 (train.py:63) y, en mi secuencia, el gradiente de GPT-2 small medía 9,76. Casi todo estaba en la tabla de tokens, que GPT-2 reutiliza como capa de salida, y el lado de la entrada se llevó 0,14. Una vez recortado medía 1,0. El recorte acorta el gradiente, pero nunca cambia su dirección.

La precisión trae su propio desvanecimiento. fp16, el formato de coma flotante de 16 bits, no puede representar un número positivo por debajo de unos 6 × 10⁻⁸, y en ese formato el 9,3 % de los valores del gradiente de GPT-2 se redondeaba a cero. Con bf16, un formato de 16 bits que conserva el rango de fp32, no se perdió ninguno.

Para fp16, el arreglo es multiplicar por un factor grande antes de la pasada hacia atrás. Cuando entrena en fp16, nanoGPT activa un GradScaler (train.py:196) cuyo factor empieza en 65 536. En mis gradientes, ese factor rescató todos menos 1743 de esos 11,6 millones de valores y empujó 11 de los más grandes por encima de 65 504, el techo de fp16. Por eso el GradScaler vuelve a dividir por el factor antes de la actualización y, cuando algo desborda, se salta esa actualización y divide el factor entre dos.

Bugs sin stack trace

En la mayoría de los programas, un bug detiene la ejecución en una línea concreta. Un gradiente equivocado no detiene nada y se limita a entrenar un modelo ligeramente distinto. He entrenado con el texto de este blog un modelo de caracteres de cuatro bloques, un GPT pequeño que lee un carácter por token. Han sido cinco ejecuciones de 1000 pasos, una correcta y cuatro con un bug distinto cada una, y todas las he puntuado con texto que el modelo no vio durante el entrenamiento. Adivinar letras solo por su frecuencia da 3,25 y la ejecución correcta ha llegado a 1,74.

Dos de los bugs se vieron en la curva. Sin zero_grad, los gradientes se fueron amontonando de un paso a otro. La pérdida bajó a 3,37 en el paso 50 y después subió hasta un 7,37 que ya era peor que una predicción uniforme entre 106 caracteres. Con los objetivos sin desplazar, al modelo se le pedía el carácter que acababa de recibir, con lo que la pérdida de entrenamiento bajó a 0,001 mientras la predicción real del carácter siguiente daba 10,7.

Los otros dos no se vieron. Con la pérdida multiplicada por 40, el bug que aparece si olvidas la división de nanoGPT, el modelo quedó tan bien como el correcto, con 1,735. AdamW, el optimizador de nanoGPT, reescala el paso de cada peso, así que un factor constante se cancela (la parte 9 explica por qué). La tabla de posiciones desconectada del grafo, a la que por eso nunca llegó gradiente, no aprendió nada y se quedó en 1,95 (un 12 % peor) con una curva que parece normal.

Para ser justos con lo que puede decir una curva, dos de los cuatro bugs los habría visto cualquiera que mirase la pérdida de entrenamiento, uno de ellos porque la pérdida salía demasiado buena. Los otros dos dejaron una curva que aceptarías, así que quien solo registra la pérdida caza la mitad.

Qué vigila el ingeniero de ML

Son cuatro hábitos, y ninguno exige saber cálculo.

Registra la norma del gradiente por bloque. Es lo más parecido a un stack trace que tiene un entrenamiento. Ahí se ve un pico, un bloque atascado cerca de cero o un total que sube poco a poco a lo largo de miles de pasos. La tabla de posiciones desconectada habría salido con gradiente cero desde el primer paso. Registra la norma antes del recorte (clip_grad_norm_ devuelve el total sin recortar), porque después el total nunca pasa del límite y el pico de cada bloque queda reducido con él.

Evalúa con una tarea en la que la pérdida de entrenamiento no pueda hacer trampa. La ejecución con los objetivos sin desplazar fue la que tuvo la mejor pérdida de entrenamiento de las cinco. Calcula una puntuación sobre texto de test, el que el modelo nunca ve al entrenar, con código que no comparta el camino de datos del bucle de entrenamiento, para que un mismo bug no pueda esconderse en ambos.

Ten claro qué bugs te esconde el optimizador. Con AdamW, multiplicar la pérdida por 40 no cambió nada. Con SGD a secas, el mismo bug equivale a una tasa de aprendizaje 40 veces más alta de la cuenta. Comprueba el escalado de la pérdida directamente, con un batch y un gradiente conocido.

Calcula la memoria de las activaciones, además de la de los pesos. GPT-2 small guarda unos 820 MB por cada secuencia de 1024 tokens para la pasada hacia atrás. Con el checkpointing de activaciones, que recalcula cada bloque durante esa pasada para no tener que guardar sus tensores, un batch de ocho pasó de 6,5 GB a 2,0 GB, a cambio de un 17 % más de tiempo en la CPU de mi portátil. La mayor parte de lo que queda son 1,65 GB de log-probabilidades, una por cada uno de los 50 257 tokens en cada posición.

Esto tampoco es trabajo de modelado. La retropropagación es una ordenación y un bucle sobre closures, y cada uno de los bugs de arriba es un error de contabilidad dentro de ese bucle o a su alrededor. Como no puede fallar haciendo ruido, el ruido lo tiene que hacer el registro.