Descenso por gradiente

Optimizadores, con foco en AdamW. Episodio uno: descenso por gradiente. Casi todos los modelos que vimos en las otras series se entrenan con AdamW, y en Alineación pasó en una sola escena. Esta serie lo arma pieza por pieza, con cada optimizador escrito a mano y comparado con PyTorch. Hoy, la pieza que está debajo de todas: dar un paso cuesta abajo.

El mapa de la serie

Diez episodios. Los dos primeros son el paso: hoy, con todos los datos; en el dos, con grupos chicos de datos, que es como se entrena de verdad. Después, tres ideas que se van sumando: recordar los pasos anteriores, que es el momentum; darle a cada peso su propio learning rate; y Adam, que junta las dos. Luego, el weight decay, y por qué en Adam hubo que desacoplarlo: eso es la W de AdamW. Más tarde, cómo cambiar el learning rate durante el entrenamiento y cómo evitar que explote. Y al final, los optimizadores que vinieron después.

Descenso → SGD minibatches → Momentum → Tasas adaptativas → Adam → Weight decay ≠ L2 → AdamW → Schedules → Estabilidad → Adafactor, Lion

Parte 1 · Un paso en una dimensión

Un solo número para mover y una pérdida con forma de parábola. Todo lo demás es esto, repetido.

Descenso por gradiente, en cinco líneas

Un único parámetro, x, que es un número suelto, y una pérdida, f de x, igual a x al cuadrado: una parábola con el mínimo en cero. Su derivada, dos x, es la pendiente de la parábola en cada punto. Empezamos en x igual a tres, con un learning rate de cero coma uno. Y la única línea que importa: a x le restamos el learning rate por la pendiente. Tres pasa a dos coma cuatro, después uno coma noventa y dos, uno coma cincuenta y tres, uno coma veintitrés. En el paso veinte, x vale cero coma cero tres cinco y la pérdida, cero coma cero cero uno. Llegamos casi al fondo.

def f(x):  return x ** 2         # la pérdida: una parábola
def df(x): return 2 * x          # su derivada: la pendiente

x, lr = 3.0, 0.1                 # punto de partida y learning rate
for t in range(1, 21):
    x = x - lr * df(x)           # un paso cuesta abajo
    if t <= 4 or t == 20:
        print(f"paso {t:2}: x = {x:.3f}   f(x) = {f(x):.3f}")

paso  1: x = 2.400   f(x) = 5.760
paso  2: x = 1.920   f(x) = 3.686
paso  3: x = 1.536   f(x) = 2.359
paso  4: x = 1.229   f(x) = 1.510
paso 20: x = 0.035   f(x) = 0.001

El paso, con números

la pendiente en el punto de partida

$\eta$: el learning rate

sacamos factor común $x_t$

Lo mismo, dibujado. La parábola, f de x igual a x al cuadrado. Su derivada, f prima, es dos x. En el dibujo, la recta que toca la curva en tres: su inclinación, seis, es la derivada. Positiva: la pérdida sube hacia la derecha, así que vamos a la izquierda. El paso: x menos eta por la derivada. Eta, la letra griega, es el learning rate: un escalar, el largo del paso. Tres menos cero coma uno por seis da dos coma cuatro. En el dibujo, el punto baja con los valores del snippet. La forma general: x sub t es el valor en el paso t. Como la derivada es dos x, sacamos factor común: x sub t más uno es uno menos dos eta, por x sub t. Con eta cero coma uno, cero coma ocho: cada paso conserva el ochenta por ciento de la distancia.

Parte 2 · Muy chico, muy grande

El learning rate es el número que más se toca al entrenar. Probemos seis valores sobre la misma parábola.

Seis learning rates, veinte pasos

La misma regla en una función, que guarda todos los x. El paso, con la derivada escrita adentro: dos x. Seis learning rates, de cero coma cero uno a uno coma uno, veinte pasos cada uno desde x igual a tres. Imprimimos también el factor uno menos dos eta. Con cero coma cero uno, después de veinte pasos sigue en dos: es lentísimo. Con cero coma uno, cero coma cero tres cinco. Con cero coma cinco, cero exacto en un solo paso. Con cero coma nueve, otra vez cero coma cero tres cinco. Con uno, vuelve a tres. Y con uno coma uno, x vale ciento quince: se fue.

def descenso(lr, x=3.0, pasos=20):
    xs = [x]
    for _ in range(pasos):
        x = x - lr * 2 * x       # f(x) = x², f'(x) = 2x
        xs.append(x)
    return xs

for lr in [0.01, 0.1, 0.5, 0.9, 1.0, 1.1]:
    xs = descenso(lr)
    print(f"lr {lr:4}: factor {1 - 2*lr:5.2f}  x final {xs[-1]:8.3f}")

lr 0.01: factor  0.98  x final    2.003
lr  0.1: factor  0.80  x final    0.035
lr  0.5: factor  0.00  x final    0.000
lr  0.9: factor -0.80  x final    0.035
lr  1.0: factor -1.00  x final    3.000
lr  1.1: factor -1.20  x final  115.013

Todo depende del factor 1 − 2η

se achica, pero muy despacio

negativo: salta al otro lado

|factor| > 1: diverge

la condición para converger en esta parábola

El factor uno menos dos eta explica las seis filas. Con eta cero coma cero uno el factor es cero coma noventa y ocho: en el dibujo, el punto casi no se mueve. Converge, pero tarda cientos de pasos. Con eta cero coma nueve el factor es menos cero coma ocho. El signo menos quiere decir que el paso se pasa del mínimo. En el dibujo, zigzag, cada vez más cerca. Con uno coma uno, menos uno coma dos: cada salto termina más lejos que el anterior. En el dibujo, arrancando en uno para que entre, rebota hacia afuera. Eso es divergir. Las barras verticales son el valor absoluto. La regla converge si el valor absoluto del factor es menor que uno, es decir, si eta está entre cero y uno. En uno exacto, rebota para siempre entre tres y menos tres.

La pérdida, paso a paso

Así se ve la pérdida, f de x, en cada uno de los veinte pasos. Con cero coma cero uno baja en una línea casi recta y muy tendida: después de veinte pasos sigue arriba de cuatro. Con cero coma uno, cae rápido y se acerca a cero. Con cero coma nueve la curva sería exactamente la misma: x salta de lado a lado, pero al cuadrado el signo no importa. Y con uno, la pérdida se queda clavada en nueve. Hay una ventana: fuera de ella, o no avanzás o explotás.

Parte 3 · Dos dimensiones

Con dos parámetros, la pérdida es un paisaje. Y aparece el problema que motiva toda la serie.

La misma regla sobre un cuenco y un valle

Ahora el parámetro es un punto, p, con dos coordenadas. El cuenco es x al cuadrado más y al cuadrado; su gradiente es dos x, dos y. El valle es cero coma uno x al cuadrado más tres y al cuadrado: muy plano en x y muy empinado en y. Su gradiente es cero coma dos x, seis y. La regla es la misma de antes, aplicada a las dos coordenadas a la vez. Veinte pasos por corrida; en el cuenco, desde dos puntos distintos. En el cuenco, cero coma uno y cero coma nueve terminan igual de cerca, a tres centésimas del centro. En el valle, con cero coma cero cinco, x casi no se movió: menos tres coma seis ocho. Con cero coma tres, menos uno coma tres. Y con cero coma tres cuatro, y se fue a tres coma dos nueve.

def g_cuenco(p): return 2 * p                          # f = x² + y²
def g_valle(p):  return np.array([0.2, 6.0]) * p       # f = 0,1x² + 3y²

def descenso(grad, p0, lr, pasos=20):
    p = np.array(p0, dtype=float); camino = [p.copy()]
    for _ in range(pasos):
        p = p - lr * grad(p)                           # la misma regla, en 2D
        camino.append(p.copy())
    return np.array(camino)
caminos = {}
for nombre, g, p0, lr in [("cuenco", g_cuenco, [-2.5, 1.5], 0.1),
                          ("cuenco", g_cuenco, [-2.5, -1.5], 0.9),
                          ("valle", g_valle, [-4.5, 1.5], 0.05),
                          ("valle", g_valle, [-4.5, 1.5], 0.3),
                          ("valle", g_valle, [-4.5, 1.5], 0.34)]:
    c = descenso(g, p0, lr)
    caminos[f"{nombre} {lr}"] = c
    print(f"{nombre:6} lr {lr:4}: ({c[-1][0]:6.3f}, {c[-1][1]:6.3f})")

cuenco lr  0.1: (-0.029,  0.017)
cuenco lr  0.9: (-0.029, -0.017)
valle  lr 0.05: (-3.681,  0.001)
valle  lr  0.3: (-1.305,  0.017)
valle  lr 0.34: (-1.100,  3.287)

Cuenco: dos learning rates

Curvas de nivel, como en un mapa de montaña: cada anillo junta los puntos de igual pérdida, y la estrella es el mínimo. Con cero coma uno, la bolita va derecho al centro: en el cuenco el gradiente apunta justo al mínimo, desde cualquier lado. Con cero coma nueve, desde abajo, cruza el centro en cada paso y salta a la esquina opuesta. Terminan igual de cerca: factor cero coma ocho para uno, menos cero coma ocho para el otro.

Por qué el valle es difícil

el gradiente: una derivada por coordenada

dos parábolas independientes

la dirección empinada pone el techo

La nabla es el gradiente, como en Álgebra lineal para LLMs, episodio nueve: una derivada parcial por coordenada, derivando por una variable con la otra quieta. Acá da cero coma dos x y seis y. En el dibujo, las dos componentes en el punto de partida: menos cero coma nueve en x y nueve en y. Cada coordenada se mueve sola, como la parábola de antes, con su propio factor. En el dibujo, con eta cero coma tres: x conserva el noventa y cuatro por ciento, y salta de lado. El problema: un solo eta para las dos. La dirección empinada exige eta menor que dos sextos, cero coma tres tres tres. En el dibujo, con cero coma tres cuatro el factor de y pasa de uno. Y con ese techo, la dirección plana avanza a paso de tortuga.

Valle: tres learning rates

Los tres caminos del snippet, veinte pasos cada uno, sobre el valle alargado. Con cero coma cero cinco, la bolita baja enseguida al fondo del valle y después casi no avanza hacia la estrella. Con cero coma tres, rebota entre las paredes y avanza un poco más: termina en menos uno coma tres. Con cero coma tres cuatro, los rebotes crecen en vez de achicarse: un poco por encima del techo y la dirección empinada diverge. Este es el problema de toda la serie. Momentum, en el episodio tres, y los learning rates por parámetro, en el cuatro, son dos maneras de resolverlo.

El techo, calculado con la hessiana

El techo de cero coma tres tres tres no hay que adivinarlo. La hessiana es la tabla de las segundas derivadas: cuánto cambia la pendiente cuando te movés. Para dos parámetros es una matriz de dos por dos, y PyTorch la calcula sola. Sus autovalores, los de Álgebra lineal para LLMs, episodio seis, son la curvatura en cada dirección principal. Acá son los números de la diagonal. El learning rate máximo es dos dividido la curvatura más grande. Y el cociente entre la curvatura más grande y la más chica se llama condition number. Curvaturas cero coma dos y seis. Learning rate máximo, cero coma tres tres tres. Condition number, treinta: una dirección es treinta veces más empinada que la otra.

def valle(p):
    return 0.1 * p[0] ** 2 + 3 * p[1] ** 2

H = torch.autograd.functional.hessian(valle, torch.zeros(2))
curv = torch.linalg.eigvalsh(H)          # curvatura en cada dirección
print("hessiana:", [[round(v, 3) for v in fila] for fila in H.tolist()])
print("curvaturas:", [round(c, 3) for c in curv.tolist()])
print(f"lr máximo = 2 / {curv.max():.1f} = {2 / curv.max():.3f}")
print(f"condition number = {curv.max() / curv.min():.1f}")

hessiana: [[0.2, 0.0], [0.0, 6.0]]
curvaturas: [0.2, 6.0]
lr máximo = 2 / 6.0 = 0.333
condition number = 30.0

Goh, 2017 — Why Momentum Really Works · fuente

0 < αλ < 2 · «All workable step-sizes, therefore, fall in the interval 0 < αλᵢ < 2» (α es nuestro η)

Esta cuenta está hecha en general en un artículo de Gabriel Goh, en la revista Distill, de dos mil diecisiete. Para cualquier pérdida cuadrática, cada dirección principal se comporta como nuestra parábola, y converge si el learning rate por la curvatura, lambda, queda entre cero y dos. Él le dice alfa al learning rate. Para él, el condition number mide la curvatura patológica: nuestro valle. Lo volvemos a abrir en el episodio tres.

Parte 4 · Lo mismo en PyTorch

torch.optim.SGD sin momentum es exactamente esta regla. Lo comprobamos número por número.

A mano contra torch.optim.SGD

SGD, stochastic gradient descent, es el nombre del optimizador más simple de PyTorch. Sin momentum, hace lo mismo que nuestra línea. Probamos en el valle, con learning rate cero coma tres. Acá los parámetros se llaman w, un vector de dos números. backward calcula el gradiente, g, y lo deja en w punto grad. step resta eta por g. Al lado, la misma cuenta a mano con NumPy, sin autograd. Y comparamos el valor absoluto de la diferencia, coordenada por coordenada. Los dos terminan en menos uno coma tres cero cinco cinco y cero coma cero uno siete tres. La diferencia, una diezmillonésima, es redondeo de float32. Lo de estocástico, en el episodio dos.

Cada línea, su término:

• perdida.backward() → $g = \nabla_w L = (0{,}2\,w_1,\; 6\,w_2)$ — backward llena w.grad
• opt.step() → $w \leftarrow w - \eta\, g$ — $\eta$ = lr = 0,3
• mano = mano - 0.3 → $w \leftarrow w - 0{,}3 \cdot (0{,}2\,w_1,\; 6\,w_2)$ — la misma cuenta, en NumPy
• diferencia máxima → $\max_i |w_i^{\text{torch}} - w_i^{\text{mano}}|$ — iguales hasta la precisión de float32

w = torch.tensor([-4.5, 1.5], requires_grad=True)   # PyTorch
opt = torch.optim.SGD([w], lr=0.3)
mano = np.array([-4.5, 1.5])                        # a mano, NumPy

for t in range(20):
    opt.zero_grad()
    perdida = 0.1 * w[0] ** 2 + 3 * w[1] ** 2
    perdida.backward()                              # w.grad = (0,2x, 6y)
    opt.step()                                      # w -= lr * w.grad
    mano = mano - 0.3 * np.array([0.2, 6.0]) * mano

print("torch:", w.detach().numpy().round(4))
print("mano: ", mano.round(4))
dif = np.abs(w.detach().numpy() - mano).max()
print(f"diferencia máxima: {dif:.1e}")

torch: [-1.3055  0.0173]
mano:  [-1.3055  0.0173]
diferencia máxima: 1.1e-07

Lo que hay que llevarse

• El gradiente dice hacia dónde sube la pérdida
• El learning rate $\eta$ decide el largo del paso
• La dirección más empinada pone el techo: $\eta < 2/\lambda_{\max}$
• Con ese techo, la dirección plana avanza lento

El gradiente apunta hacia donde la pérdida sube más rápido; el paso va en contra. El learning rate es el largo del paso, y hay una ventana: muy chico no avanza, muy grande diverge. El techo lo pone la curvatura más grande: dos dividido lambda máximo. Y si las curvaturas son muy distintas, como en el valle, el mismo eta que es seguro para la dirección empinada es lentísimo para la plana. Esa tensión es el hilo de la serie, y la mirada general está en Álgebra, episodio diez.

Resumiendo

• Descenso por gradiente: $x \leftarrow x - \eta\, f'(x)$
• En $x^2$: cada paso multiplica por $1 - 2\eta$
• Converge si $|1 - 2\eta| < 1$; si no, zigzag que crece
• Valle: curvaturas 0,2 y 6, techo $\eta < 0{,}333$
• torch.optim.SGD sin momentum: iguales a mano
• Próximo: minibatches y gradientes con ruido

Papers del episodio

• Goh, 2017 — Why Momentum Really Works