AdamW

Optimizadores, con foco en AdamW. Episodio siete: AdamW, Adam con decoupled weight decay. En el episodio cinco armamos Adam, adaptive moment estimation, y en el seis vimos que sumarle L2 al gradiente no es lo mismo que achicar los pesos. Hoy escribimos AdamW a mano, línea por línea, lo comparamos con el de PyTorch hasta la última cifra, y vemos en un paisaje con dos mínimos dónde termina cada uno.

Un paso de AdamW

El mapa de un paso, antes del código. Primero, los pesos, theta, se multiplican por uno menos eta lambda: el weight decay, sin mirar el gradiente. Después, el gradiente, g, actualiza dos promedios móviles: m, el del gradiente, y v, el del gradiente al cuadrado. Los dos se corrigen, porque arrancan en cero. Y el paso: eta por m sombrero sobre la raíz de v sombrero más épsilon. AdamW solo agrega la primera caja.

θ → θ · (1 − ηλ) → g → m → v → m̂ y v̂ → η · m̂ / (√v̂ + ε) → θ nuevo

Parte 1 · Un paso, con números

Dos pesos, un gradiente, y cada línea impresa.

Las cinco líneas de un paso

Theta vale uno y menos dos; el gradiente, cero coma cinco y menos cuatro. Eta y lambda valen cero coma uno; betas y épsilon, los del episodio cinco; es el paso t igual a uno. El decay: multiplicar por cero coma noventa y nueve. m: beta uno por el m anterior, que es cero, más uno menos beta uno por g. v: lo mismo con g al cuadrado. La corrección de sesgo divide por uno menos beta a la t. Y el paso. El decay deja cero coma noventa y nueve y menos uno coma noventa y ocho. m da cero coma cero cinco y menos cero coma cuatro. Corregidos, vuelven a ser el gradiente y su cuadrado. Y theta termina en cero coma ochenta y nueve y menos uno coma ochenta y ocho: los dos se movieron exactamente cero coma uno.

th = torch.tensor([1.0, -2.0])                  # θ: dos pesos
g = torch.tensor([0.5, -4.0])                   # g: su gradiente en t = 1
eta, lam, b1, b2, eps, t = 0.1, 0.1, 0.9, 0.999, 1e-8, 1

th = th * (1 - eta * lam)                       # 1. weight decay
print("1. decay:", r(th))
m = b1 * 0 + (1 - b1) * g                       # 2. primer momento
v = b2 * 0 + (1 - b2) * g ** 2                  # 3. segundo momento
print("2-3. m:", r(m), " v:", r(v, 5))
mh, vh = m / (1 - b1 ** t), v / (1 - b2 ** t)  # 4. corrección de sesgo
print("4. m̂:", r(mh), " v̂:", r(vh))
th = th - eta * mh / (vh.sqrt() + eps)          # 5. el paso
print("5. θ:", r(th))

1. decay: [0.99, -1.98]
2-3. m: [0.05, -0.4]  v: [0.00025, 0.016]
4. m̂: [0.5, -4.0]  v̂: [0.25, 16.0]
5. θ: [0.89, -1.88]

El mismo paso, término a término

1. el decay, antes de todo

2. $\beta_1 = 0{,}9$; $v_1$ igual, con $g^2$ y $\beta_2$

3. la corrección deshace el arranque en cero

4. el primer paso es η por el signo de g

El decay multiplica cada peso por cero coma noventa y nueve. En el dibujo, los dos pesos pierden el uno por ciento, sin mirar el gradiente. m sub uno es beta uno por m sub cero, que vale cero, más uno menos beta uno por g. En el dibujo, la suma con pesos cero coma nueve y cero coma uno: solo entra el diez por ciento del gradiente. v sub uno, igual, con g al cuadrado. La corrección divide por uno menos beta uno a la t: acá, cero coma uno. En el dibujo, menos cero coma cuatro vuelve a ser menos cuatro, el gradiente entero. El paso: m sombrero sobre la raíz de v sombrero. La raíz de g al cuadrado es el valor absoluto de g, así que el cociente es el signo. En el dibujo, gradientes ocho veces distintos y pasos idénticos.

Parte 2 · La clase completa

Las mismas líneas, ahora para cualquier lista de tensores y cualquier cantidad de pasos.

AdamW a mano

Una clase con la interfaz y los valores por defecto de torch punto optim. Guarda un m y un v por tensor, del mismo tamaño que el tensor, y un contador t. En step, para cada tensor: primero el decay, en el lugar, con mul guion bajo. Es lo que AdamW agrega. Después m, beta uno por m más uno menos beta uno por el gradiente. Después v, igual con el gradiente al cuadrado, elemento a elemento. Y el paso: restar eta por m sombrero sobre la raíz de v sombrero más épsilon. El no grad de arriba evita que estas cuentas entren al grafo.

Cada línea, su término:

• p.mul_(1 - self.lr * self.wd) → $\theta \leftarrow (1 - \eta\lambda)\,\theta$ — decay, antes de tocar m y v
• m.mul_(self.b1) → $m \leftarrow \beta_1 m + (1 - \beta_1)\,g$ — promedio móvil del gradiente
• v.mul_(self.b2) → $v \leftarrow \beta_2 v + (1 - \beta_2)\,g^2$ — promedio móvil del cuadrado, elemento a elemento
• p.sub_(self.lr * mh → $\theta \leftarrow \theta - \eta\,\frac{\hat m}{\sqrt{\hat v} + \epsilon}$ — $\hat m, \hat v$: corregidos por $1 - \beta^t$

class AdamWMano:
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999),
                 eps=1e-8, weight_decay=1e-2):
        self.p, self.lr, self.wd, self.eps = list(params), lr, weight_decay, eps
        self.b1, self.b2, self.t = *betas, 0
        self.m = [torch.zeros_like(p) for p in self.p]
        self.v = [torch.zeros_like(p) for p in self.p]

    @torch.no_grad()
    def step(self):
        self.t += 1
        for p, m, v in zip(self.p, self.m, self.v):
            p.mul_(1 - self.lr * self.wd)                       # decay, aparte
            m.mul_(self.b1).add_((1 - self.b1) * p.grad)        # m
            v.mul_(self.b2).add_((1 - self.b2) * p.grad ** 2)   # v
            mh = m / (1 - self.b1 ** self.t)
            vh = v / (1 - self.b2 ** self.t)
            p.sub_(self.lr * mh / (vh.sqrt() + self.eps))       # paso

A mano contra torch.optim.AdamW

La prueba. Dos copias de una capa lineal de cinco entradas y tres salidas. La segunda arranca con los mismos pesos que la primera. La misma configuración para los dos optimizadores: learning rate cero coma cero uno, lambda cero coma uno. Uno es nuestra clase y el otro, el de PyTorch. Doscientos pasos de una pérdida cuadrática sobre datos al azar. Y comparamos peso por peso. Las dos pérdidas terminan en uno coma uno siete ocho cuatro uno. La diferencia máxima entre pesos es tres por diez a la menos ocho: redondeo de float32. Al final, los defaults de PyTorch: learning rate diez a la menos tres, épsilon diez a la menos ocho, weight decay cero coma cero uno.

red_a = torch.nn.Linear(5, 3)
red_b = torch.nn.Linear(5, 3)
red_b.load_state_dict(red_a.state_dict())                 # mismos pesos
cfg = dict(lr=0.01, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.1)
mano = AdamWMano(red_a.parameters(), **cfg)
ref = torch.optim.AdamW(red_b.parameters(), **cfg)
for paso in range(200):
    for red, opt in [(red_a, mano), (red_b, ref)]:
        for p in red.parameters(): p.grad = None
        ((red(X) - Y) ** 2).mean().backward()
        opt.step()
dif = max((a - b).abs().max().item()
          for a, b in zip(red_a.parameters(), red_b.parameters()))
print("pérdida a mano:", round(((red_a(X) - Y) ** 2).mean().item(), 5))
print("pérdida torch :", round(((red_b(X) - Y) ** 2).mean().item(), 5))
print(f"máx |diferencia| en los pesos: {dif:.1e}")
print("defaults de torch:", {k: torch.optim.AdamW([torch.zeros(1)]).defaults[k]
                              for k in ("lr", "eps", "weight_decay")})

pérdida a mano: 1.17841
pérdida torch : 1.17841
máx |diferencia| en los pesos: 3.0e-08
defaults de torch: {'lr': 0.001, 'eps': 1e-08, 'weight_decay': 0.01}

Loshchilov y Hutter, 2017 — Decoupled Weight Decay Regularization · arXiv:1711.05101

Algorithm 2, línea 12 · «We decouple weight decay and loss-based gradient updates in Adam as shown in line 12 of Algorithm 2; this gives rise to our variant of Adam with decoupled weight decay (AdamW).»

Así lo escriben Loshchilov y Hutter. En la línea doce de su algoritmo dos, el paso es theta menos eta sub t por un paréntesis: alfa por m sombrero sobre la raíz de v sombrero más épsilon, más lambda theta. Ojo con los nombres: para ellos, alfa es el learning rate, y eta sub t es un multiplicador que cambia con el tiempo, el schedule, que es el episodio ocho. Con eso, el decay del paper no se multiplica por alfa.

El λ del paper y el de PyTorch

$\alpha$: learning rate; $\eta_t$: schedule

$\gamma$: el lr de torch

en torch, subir el lr también sube el decay

Esto importa al copiar hiperparámetros de un paper. En el paper, alfa es el learning rate, por ejemplo cero coma cero cero uno, y eta sub t es el schedule, que acá vale uno. Lambda solo se multiplica por eta sub t. En el dibujo, con lambda cero coma cero cero cero uno, se pierde una diezmilésima por paso. En PyTorch hay un único learning rate, gamma, que ya incluye el schedule. Nuestra clase lo confirmó: el decay es gamma por weight decay. En el dibujo, weight decay cero coma uno en torch da el mismo decay por paso. Así que lambda del paper es alfa por el lambda de torch. Consecuencia práctica: en torch, si triplicás el learning rate, triplicás también el decay. En el dibujo, de una a tres diezmilésimas por paso.

Parte 3 · Dos mínimos

Adam, Adam + L2 y AdamW sobre el mismo paisaje, desde el mismo punto, con el mismo λ.

Tres optimizadores, un paisaje

La pérdida tiene dos mínimos: x al cuadrado menos uno, al cuadrado, más cero coma cinco y al cuadrado, más cero coma dos x, que hunde un poco el de la izquierda. Arrancamos en uno coma ocho, uno coma cinco, con learning rate cero coma cero cinco. Además, leemos v del estado del optimizador y calculamos la raíz de v sombrero de x. Adam con weight decay cero coma cinco: L2, como vimos en el seis. Y AdamW con el mismo número. Adam, sin castigo, para en cero coma nueve siete cuatro. Adam con L2, en cero coma nueve cero cinco. AdamW, en el paso doscientos, en cero coma ocho dos cuatro; en el ochocientos, en cero coma nueve cero tres. La raíz de v sombrero bajó de dos coma cero nueve a uno coma cero cuatro.

def f(p):                                       # dos mínimos: (x² − 1)² + 0,5y² + 0,2x
    return (p[0] ** 2 - 1) ** 2 + 0.5 * p[1] ** 2 + 0.2 * p[0]

def camino(Clase, wd, pasos):
    p = torch.tensor([1.8, 1.5], requires_grad=True)
    opt = Clase([p], lr=0.05, weight_decay=wd)
    puntos = [p.tolist()]
    for _ in range(pasos):
        opt.zero_grad(); f(p).backward(); opt.step()
        puntos.append(p.tolist())
    est = opt.state[p]                          # √v̂ de x: la escala que recuerda Adam
    rv = (est["exp_avg_sq"][0] / (1 - 0.999 ** est["step"])).sqrt().item()
    return np.array(puntos), rv
cam = {"Adam": camino(torch.optim.Adam, 0.0, 200),
       "Adam + L2": camino(torch.optim.Adam, 0.5, 200),
       "AdamW": camino(torch.optim.AdamW, 0.5, 200)}
for k, (c, rv) in cam.items():
    print(f"{k:9} paso 200: x = {c[-1, 0]:.3f}")
c, rv800 = camino(torch.optim.AdamW, 0.5, 800)
print(f"AdamW     paso 800: x = {c[-1, 0]:.3f}")
print(f"AdamW √v̂ de x: paso 200 {cam['AdamW'][1]:.2f}, paso 800 {rv800:.2f}")

Adam      paso 200: x = 0.974
Adam + L2 paso 200: x = 0.905
AdamW     paso 200: x = 0.824
AdamW     paso 800: x = 0.903
AdamW √v̂ de x: paso 200 2.09, paso 800 1.04

Doscientos pasos en dos mínimos

Los tres caminos. La estrella marca el mínimo de la derecha, en x igual a cero coma nueve siete cuatro. Adam baja en diagonal, porque sus primeros pasos tienen el mismo largo en x y en y, y se queda en la estrella. Adam con L2 hace casi lo mismo, pero se detiene antes: en cero coma nueve cero cinco, un poco más cerca del origen, que es adonde empuja el castigo. AdamW se adelanta hacia el origen y a los doscientos pasos todavía está en cero coma ocho dos cuatro, del lado de la loma. ¿Por qué tres lugares distintos?

La coordenada x, paso a paso

Solo la coordenada x, en los mismos doscientos pasos. Adam baja y se planta en cero coma nueve siete cuatro. Adam con L2, en cero coma nueve cero cinco. AdamW cae hasta cero coma cinco, en el paso treinta y nueve, y después sube despacio: todavía no llegó.

Dónde se detiene cada uno

para donde el gradiente con castigo se anula

el decay contra un paso dividido por √v̂

v se olvida despacio: β₂ = 0,999

Adam con L2 se detiene donde el gradiente que recibe, f prima más lambda x, vale cero. En el dibujo, en cero coma nueve cero cinco la derivada es menos cero coma cuatro cinco cinco y el castigo, más cero coma cuatro cinco: se cancelan. Ese punto no depende de v. AdamW se detiene donde el decay, lambda x, empata con el paso de Adam, que es el gradiente dividido la raíz de v sombrero. En el dibujo, en el paso doscientos: cero coma cuatro uno dos contra cero coma cuatro uno, empatados. Pero v recuerda los gradientes enormes del arranque, y con beta dos cero coma nueve nueve nueve los olvida despacio. Cuando la raíz de v baja a uno coma cero cuatro, el punto de equilibrio se mueve a cero coma nueve cero tres. En el dibujo, las cuatro posiciones finales.

Parte 4 · AdamW en la práctica

Qué parámetros llevan decay, y qué conviene saber antes de copiar un número.

Sin decay en sesgos ni LayerNorm

Es una costumbre muy extendida aplicar el decay solo a las matrices. Las matrices son los parámetros de dos dimensiones o más. Los vectores, de una dimensión, son los sesgos y las escalas de LayerNorm, la normalización de cada capa. torch punto optim acepta grupos de parámetros, cada uno con su weight decay: cero coma uno para las matrices, cero para el resto. Para ver solo el decay, damos cien pasos con gradiente cero: Adam no se mueve, el decay sí. Las dos matrices conservan cero coma nueve cero cuatro ocho, que es cero coma nueve nueve nueve a la cien. Los sesgos y la escala de LayerNorm, uno: intactos.

modelo = torch.nn.Sequential(torch.nn.Linear(4, 8), torch.nn.LayerNorm(8),
                             torch.nn.Linear(8, 2))
con = [p for p in modelo.parameters() if p.dim() >= 2]    # matrices
sin = [p for p in modelo.parameters() if p.dim() < 2]     # sesgos y LayerNorm
opt = torch.optim.AdamW([{"params": con, "weight_decay": 0.1},
                         {"params": sin, "weight_decay": 0.0}], lr=0.01)
antes = [p.detach().norm().item() for p in modelo.parameters()]
for _ in range(100):
    for p in modelo.parameters(): p.grad = torch.zeros_like(p)
    opt.step()                                            # solo actúa el decay
for (n, p), a in zip(modelo.named_parameters(), antes):
    queda = f"{p.norm().item() / a:.4f}" if a else "era 0"
    print(f"{n:9} {tuple(p.shape)!s:7} queda {queda}")

0.weight  (8, 4)  queda 0.9048
0.bias    (8,)    queda 1.0000
1.weight  (8,)    queda 1.0000
1.bias    (8,)    queda era 0
2.weight  (2, 8)  queda 0.9048
2.bias    (2,)    queda 1.0000

Qué hace weight_decay en cada clase

Un resumen de los tres episodios en una tabla. En SGD, stochastic gradient descent, weight decay en torch es L2, que en SGD es lo mismo que el decay: cada peso se achica por uno menos eta lambda. En Adam, el mismo parámetro es L2, y como pasa por la raíz de v, cada peso recibe un castigo distinto. En AdamW, es el decay desacoplado: uno menos eta lambda, igual para todos, con eta el learning rate de torch.

Antes de copiar hiperparámetros

• Defaults de torch: lr $10^{-3}$, weight_decay 0,01
• En torch el decay por paso es lr · weight_decay
• El λ de un paper puede no incluir el lr
• Matrices con decay; sesgos y LayerNorm sin

Cuatro cosas para tener a mano. Los defaults de AdamW en PyTorch, que imprimió el snippet: learning rate diez a la menos tres y weight decay cero coma cero uno. El decay por paso es el producto del learning rate y el weight decay: si cambiás uno, cambia el efecto del otro. El paper original multiplica lambda solo por el schedule. Antes de copiar un lambda, mirá en qué convención está escrito. Y la costumbre de dejar fuera del decay a los sesgos y a LayerNorm, con grupos de parámetros.

Resumiendo

• AdamW = Adam + achicar θ por $1 - \eta\lambda$ antes del paso
• Primer paso de Adam: $\eta$ por el signo de $g$
• A mano contra torch.optim.AdamW: diferencia $3 \cdot 10^{-8}$
• Paper: decay por $\eta_t$; torch: por lr · weight_decay
• Adam + L2 para en 0,905; AdamW depende de $\sqrt{\hat v}$
• Próximo: schedules, warmup y coseno

Papers del episodio

• Loshchilov y Hutter, 2017 — Decoupled Weight Decay Regularization