Optimización

Álgebra lineal para LLMs. Episodio diez, el último: optimización. Ya sabemos calcular gradientes; ahora, qué se hace con ellos. Bajar por la pérdida, hacerlo bien con Adam, adaptive moment estimation, y optimizar con una restricción, que es lo que hay detrás de la política óptima de DPO, direct preference optimization.

Las fórmulas que nos traen

Alineación, ep. 2: el paso de AdamW

Alineación, ep. 5: la política óptima de DPO

Dos fórmulas de Alineación. La del episodio dos: el paso de AdamW, Adam con weight decay desacoplado. Theta son todos los pesos; eta, el learning rate; m sombrero y v sombrero, dos promedios del gradiente; épsilon, un número chiquito; y lambda por theta, el weight decay. En el dibujo, la idea de base: bajar por la pendiente. Y la del episodio cinco: la política óptima, pi estrella, es la referencia por e a la recompensa sobre beta, normalizada por Z. En el dibujo, las cinco respuestas de aquel episodio. Hoy vemos de dónde sale cada pieza.

Parte 1 · Bajar por el gradiente

Menos el gradiente apunta cuesta abajo. El paso más simple: moverse un poquito en esa dirección, y repetir.

Descenso por gradiente en un valle

Una pérdida de dos parámetros con forma de valle: cero coma uno x al cuadrado más tres y al cuadrado. Muy empinada en y, casi plana en x. La pérdida. Su gradiente: cero coma dos x, seis y. Probamos tres learning rates, eta, desde el mismo punto, menos cuatro coma cinco, uno coma cinco. Treinta pasos de la regla: restar eta por el gradiente. Con eta cero coma uno, y llega a cero pero x se queda en menos dos coma cuatro: pérdida cero coma seis. Con cero coma tres, pérdida cero coma cero cinco. Con cero coma treinta y cuatro, apenas más, y se va a cuatro coma ocho: diverge. La pérdida, setenta y uno.

def perdida(p):                       # el valle: 0,1·x² + 3·y²
    return 0.1 * p[0]**2 + 3 * p[1]**2

def grad(p):                          # (0,2·x, 6·y)
    return np.array([0.2 * p[0], 6 * p[1]])

for eta in [0.1, 0.3, 0.34]:          # eta: learning rate
    p = np.array([-4.5, 1.5])
    camino = [p]
    for paso in range(30):
        p = p - eta * grad(p)         # descenso por gradiente
        camino.append(p)
    print(f"eta={eta}: p=({p[0]:+.3f}, {p[1]:+.3f})  pérdida={perdida(p):.4f}")
    caminos[eta] = camino

eta=0.1: p=(-2.455, +0.000)  pérdida=0.6025
eta=0.3: p=(-0.703, +0.002)  pérdida=0.0495
eta=0.34: p=(-0.544, +4.865)  pérdida=71.0371

Por qué hay un learning rate máximo

$\eta$: learning rate · $t$: número de paso

la dirección empinada del valle

0,3 converge zigzagueando · 0,34 diverge

La regla general: theta en el paso t más uno es theta en el paso t menos eta por el gradiente. En el dibujo, x al cuadrado con eta cero coma tres: tres, uno coma dos, cero coma cuarenta y ocho. Cada paso achica x por el mismo factor. En la dirección y del valle, el gradiente es seis y, así que cada paso multiplica y por uno menos seis eta. En el dibujo, x al cuadrado con un eta demasiado grande: el factor es negativo y mayor que uno en tamaño, y cada salto cruza el mínimo y cae más lejos. Converge solo si ese factor tiene tamaño menor que uno: eta menor que dos sextos. Por eso cero coma tres zigzaguea pero llega, y cero coma treinta y cuatro explota.

El valle: η chico, η grande, η de más

Los tres caminos de verdad, sobre las curvas de nivel del valle. Con eta cero coma uno: baja enseguida al fondo del valle y después camina despacio a lo largo, porque en x el gradiente es chiquito. Con cero coma tres: zigzaguea de pared a pared, pero avanza más en x. Con cero coma treinta y cuatro, cada zigzag es más ancho que el anterior, hasta salirse. El problema de fondo: un único eta para dos direcciones con curvaturas treinta veces distintas.

Parte 2 · Momentum y Adam

Dos ideas para el valle: acumular velocidad en la dirección que se repite, y darle a cada parámetro su propio tamaño de paso.

Adam, en una línea

Adam, a mano, sobre el mismo valle. Dos coeficientes, beta uno y beta dos, el learning rate, y épsilon, para no dividir por cero. m es un promedio móvil del gradiente: noventa por ciento de lo que traía y diez por ciento del gradiente nuevo. v es lo mismo con el gradiente al cuadrado. El paso, en una línea: m corregido dividido por la raíz de v corregido, por eta. Y lo comparamos con el Adam de PyTorch. En el primer paso, el gradiente es menos cero coma nueve en x y nueve en y, diez veces más grande; pero el paso es cero coma uno en las dos coordenadas. Después de sesenta pasos, pérdida cero coma cero uno cero seis. Y coincide con torch punto optim punto Adam.

p = torch.tensor([-4.5, 1.5])
m, v = torch.zeros(2), torch.zeros(2)
b1, b2, eta, eps = 0.9, 0.999, 0.1, 1e-8
for t in range(1, 61):
    g = grad(p)
    m = b1 * m + (1 - b1) * g             # promedio móvil de g
    v = b2 * v + (1 - b2) * g**2          # promedio móvil de g²
    paso = eta * (m / (1 - b1**t)) / ((v / (1 - b2**t)).sqrt() + eps)
    p = p - paso                          # Adam, en una línea
    if t == 1:
        print("gradiente 1:", [round(x, 2) for x in g.tolist()],
              " paso 1:", [round(x, 3) for x in paso.tolist()])
print(f"a mano, 60 pasos: pérdida {perdida(p):.4f}")

q = torch.tensor([-4.5, 1.5], requires_grad=True)
opt = torch.optim.Adam([q], lr=0.1)
for _ in range(60):
    opt.zero_grad(); perdida(q).backward(); opt.step()
print("igual a torch.optim.Adam:", torch.allclose(p, q.detach(), atol=1e-5))

gradiente 1: [-0.9, 9.0]  paso 1: [-0.1, 0.1]
a mano, 60 pasos: pérdida 0.0106
igual a torch.optim.Adam: True

Momentum y Adam, término a término

momentum: $u$ suma los gradientes pasados, con peso $\beta^k$

promedios móviles exponenciales de g y de g²

cada parámetro dividido por su propia escala

Momentum: u es una velocidad, beta por la anterior más el gradiente nuevo. Es la suma de los gradientes pasados con pesos beta a la k: en el dibujo, uno, cero coma nueve, cero coma ochenta y uno. Lo que se repite se acumula; el zigzag se cancela. Adam guarda dos promedios móviles exponenciales: m, del gradiente, y v, del gradiente al cuadrado, elemento a elemento. En el dibujo, los pesos de m: los mismos beta a la k, por cero coma uno. Y el paso divide m por la raíz de v. Si un parámetro tiene gradientes grandes, la raíz de v es grande y el paso se achica; si son chicos, se agranda. En el dibujo, el primer paso: gradientes diez veces distintos, pasos iguales. Los uno menos beta a la t corrigen el arranque en cero.

Tres optimizadores, sesenta pasos

Los tres, sesenta pasos desde el mismo punto, con la pérdida final en la leyenda. Descenso por gradiente con eta cero coma uno: se queda en cero coma uno ocho. Momentum: se pasa del fondo, vuelve, pero acumula velocidad a lo largo del valle y llega a cero coma cero cero siete. Adam arranca en diagonal, porque mueve las dos coordenadas lo mismo, y llega a cero coma cero uno. En un LLM no hay dos direcciones sino miles de millones, con curvaturas muy distintas: por eso Adam es el default.

Parte 3 · Gradient clipping y AdamW

Dos detalles que están en todo entrenamiento de un LLM: ponerle techo al gradiente, y achicar los pesos sin mezclarlo con Adam.

clip_grad_norm_ y weight decay desacoplado

Mil pesos con una pérdida que da gradientes enormes. clip grad norm escala el gradiente para que su norma no pase de uno: la misma dirección, más corto. Eso es gradient clipping. AdamW con weight decay cero coma uno. Y Adam con el weight decay viejo, que lo suma al gradiente. Les damos gradiente cero y un paso. La norma del gradiente era seiscientos cincuenta mil; después del recorte, uno. AdamW deja los pesos en cero coma noventa y nueve: los multiplica por uno menos eta por lambda, aparte. Adam con L2 los deja en cero coma nueve: el decay entró al gradiente, Adam lo normalizó y lo convirtió en un paso entero. Eso es lo que AdamW, Adam con decoupled weight decay, separa.

w = torch.nn.Parameter(torch.randn(1000))
(100 * w).pow(2).sum().backward()          # gradientes enormes
antes = torch.nn.utils.clip_grad_norm_([w], max_norm=1.0)
print(f"norma antes: {antes.item():,.0f}   después: {w.grad.norm().item():.3f}")

a = torch.nn.Parameter(torch.ones(3))      # AdamW: decay desacoplado
opt_w = torch.optim.AdamW([a], lr=0.1, weight_decay=0.1)
b = torch.nn.Parameter(torch.ones(3))      # Adam con weight_decay (L2)
opt_l2 = torch.optim.Adam([b], lr=0.1, weight_decay=0.1)

a.grad, b.grad = torch.zeros(3), torch.zeros(3)   # gradiente cero
opt_w.step(); opt_l2.step()
print("AdamW:", [round(x, 3) for x in a.tolist()])
print("Adam + L2:", [round(x, 3) for x in b.tolist()])

norma antes: 650,536   después: 1.000
AdamW: [0.99, 0.99, 0.99]
Adam + L2: [0.9, 0.9, 0.9]

Kingma y Ba, 2014 — Adam: A Method for Stochastic Optimization · arXiv:1412.6980

β₁ = 0,9 · β₂ = 0,999 · «Good default settings for the tested machine learning problems are α = 0.001, β₁ = 0.9, β₂ = 0.999 and ε = 10⁻⁸.»

Los números de Adam tienen origen. Kingma y Ba presentaron Adam en dos mil catorce, y en el algoritmo escriben los valores por defecto: learning rate cero coma cero cero uno, beta uno cero coma nueve, beta dos cero coma nueve nueve nueve y épsilon diez a la menos ocho. Siguen siendo el punto de partida de casi todo entrenamiento.

Parte 4 · Optimizar con una restricción

A veces no se puede ir a cualquier lado: una distribución tiene que sumar uno. Para eso está el multiplicador de Lagrange.

Lagrange en un círculo

Un ejemplo de dos dimensiones. La restricción: solo valen los puntos del círculo de radio uno; recorremos tres mil seiscientos. Queremos el punto del círculo donde x más y es máximo. En ese punto calculamos el gradiente de f, uno, uno, y el de la restricción, dos x, dos y. Y dividimos uno por el otro, coordenada a coordenada. El mejor punto es cero coma siete cero siete en las dos coordenadas, con f igual a uno coma cuatro uno cuatro. Los dos gradientes son paralelos: uno es cero coma siete cero siete veces el otro, en las dos coordenadas. Ese número es lambda, el multiplicador de Lagrange.

Cada línea, su término:

• x, y = np.cos(ang), np.sin(ang) → $g(x, y) = x^2 + y^2 = 1$ — la restricción: solo puntos del círculo
• i = np.argmax(x + y) → $\max\; f(x, y) = x + y$ — lo que queremos maximizar
• grad_g = np.array → $\nabla f = (1, 1), \quad \nabla g = (2x, 2y)$ — los dos gradientes en el óptimo
• lam = grad_f / grad_g → $\nabla f = \lambda\, \nabla g$ — $\lambda$: el multiplicador de Lagrange

# maximizar f(x, y) = x + y  con la restricción x² + y² = 1
ang = np.linspace(0, 2 * np.pi, 3601)
x, y = np.cos(ang), np.sin(ang)          # todos los puntos del círculo
i = np.argmax(x + y)
x0, y0 = x[i], y[i]
print(f"mejor punto: ({x0:.3f}, {y0:.3f})   f = {x0 + y0:.3f}")

grad_f = np.array([1.0, 1.0])            # ∇(x + y)
grad_g = np.array([2 * x0, 2 * y0])      # ∇(x² + y²)
lam = grad_f / grad_g                    # ∇f = λ ∇g, coordenada a coordenada
print("∇f =", grad_f, "  ∇g =", grad_g.round(3))
print("λ =", lam.round(3), "  1/√2 =", round(1 / np.sqrt(2), 3))

mejor punto: (0.707, 0.707)   f = 1.414
∇f = [1. 1.]   ∇g = [1.414 1.414]
λ = [0.707 0.707]   1/√2 = 0.707

En el óptimo, los gradientes son paralelos

El círculo es la restricción. En el punto uno, cero, el gradiente de la restricción apunta hacia afuera. El de f apunta en diagonal: tiene una parte a lo largo del círculo, así que moviéndose por el círculo, f todavía sube. No es el óptimo. En el óptimo, el gradiente de f apunta hacia afuera, igual que el de la restricción: ya no hay forma de subir sin salirse. La recta x más y igual a uno coma cuatro uno, la curva de nivel de f, apenas toca el círculo ahí. Esa es la condición: gradiente de f igual a lambda por gradiente de g.

El Lagrangiano de π*

objetivo de Alineación ep. 5 + la restricción $\sum \pi = 1$

derivada de π log π: log π + 1

$\lambda$ se elige para que sume 1: $e^{\lambda/\beta + 1} = Z$

Ahora la política óptima de Alineación, episodio cinco, con esta herramienta. El objetivo: recompensa esperada menos beta por la KL, de Kullback-Leibler, contra la referencia. Le restamos lambda por la restricción: las probabilidades tienen que sumar uno. En el dibujo, las cinco de pi estrella, que suman uno. Derivamos respecto de cada pi sub y. La derivada de pi por log de pi es log de pi más uno. Igualamos a cero: en el óptimo, cada derivada vale lambda. En el dibujo, las cinco valen cero coma dos cinco uno. Despejamos con la exponencial: pi estrella es la referencia por e a la r sobre beta, por una constante. Esa constante la fija lambda para que todo sume uno, y es uno sobre Z. Es un softmax con temperatura beta, el del episodio siete.

π*, comprobada

Con los números de aquel episodio. Cinco respuestas: la referencia y sus recompensas, beta uno. Pi estrella: referencia por e a la r sobre beta, normalizada. La derivada del Lagrangiano, sin el lambda, para cada respuesta. El lambda que predice la cuenta: beta por log Z menos uno. Y pi estrella escrita como softmax. Pi estrella da once, veintitrés, treinta y dos, cinco y veintinueve por ciento. Las cinco derivadas valen cero coma dos cinco uno, todas iguales, y beta por log Z menos uno da cero coma dos cinco uno: es lambda. Y el softmax de log pi ref más r da exactamente pi estrella.

pi_ref = torch.tensor([0.40, 0.30, 0.15, 0.10, 0.05])  # Alineación ep. 5
r = torch.tensor([0.0, 1.0, 2.0, 0.5, 3.0])
beta = 1.0

pi = pi_ref * torch.exp(r / beta)          # π_ref · e^(r/β)
Z = pi.sum()
pi = pi / Z                                # normalizar: Σ π = 1
print("π* =", r3(pi))

# derivada del Lagrangiano respecto de cada π_y (sin el -λ)
d = r - beta * (torch.log(pi / pi_ref) + 1)
print("r - β(log π*/π_ref + 1) =", r3(d))
print("β(log Z - 1) =", round((beta * (Z.log() - 1)).item(), 3))
print("softmax(log π_ref + r/β) =", r3(torch.softmax(pi_ref.log() + r / beta, 0)))

π* = [0.115, 0.233, 0.317, 0.047, 0.288]
r - β(log π*/π_ref + 1) = [0.251, 0.251, 0.251, 0.251, 0.251]
β(log Z - 1) = 0.251
softmax(log π_ref + r/β) = [0.115, 0.233, 0.317, 0.047, 0.288]

Volvemos a las fórmulas de Alineación

paso normalizado por parámetro + weight decay aparte

Lagrange: Z es la normalización que impone λ

Las dos fórmulas, leídas enteras. El paso de AdamW: m sombrero sobre la raíz de v sombrero es Adam, un paso normalizado para cada parámetro, con momentum adentro. Lambda por theta es el weight decay, aplicado aparte, sin pasar por esa normalización. Y eta multiplica todo. Y pi estrella: el máximo de la recompensa menos la KL, con la restricción de sumar uno. Z es lo que impone lambda. En el dibujo, la respuesta con recompensa tres pasa de cinco a veintinueve por ciento.

Resumiendo

• Descenso por gradiente: $\theta \leftarrow \theta - \eta \nabla \mathcal{L}$; converge si $\eta < 2/\text{curvatura}$
• Momentum acumula; Adam divide por $\sqrt{v}$: un paso a medida de cada parámetro
• Gradient clipping limita la norma; AdamW aplica el weight decay aparte
• Lagrange: en el óptimo con restricción, $\nabla f = \lambda \nabla g$
• $\pi^* = \pi_{\text{ref}} e^{r/\beta} / Z$: un softmax con temperatura $\beta$
• Seguí con: «Optimizadores, con foco en AdamW» y «LoRA»

Papers del episodio

• Kingma y Ba, 2014 — Adam: A Method for Stochastic Optimization