Weight decay no es L2

Optimizadores, con foco en AdamW. Episodio seis: weight decay no es L2. Hay dos maneras de pedirle a una red que mantenga los pesos chicos. Con SGD, stochastic gradient descent, dan exactamente lo mismo, y por eso durante años se usaron como sinónimos. Con Adam, adaptive moment estimation, dan cosas distintas. Hoy lo demostramos con números, y eso prepara el episodio siete: AdamW, Adam con decoupled weight decay.

Parte 1 · Dos formas de achicar los pesos

Una suma un castigo a la pérdida. La otra multiplica los pesos por un número un poco menor que uno.

Para qué achicar los pesos

• Pesos grandes: la red memoriza (overfitting)
• L2: sumarle $\tfrac{\lambda}{2}\|\theta\|^2$ a la pérdida
• Weight decay: multiplicar $\theta$ por $(1 - \eta\lambda)$
• Mismo objetivo, distinto lugar en el código

Por qué querríamos pesos chicos. Una red con pesos enormes puede ajustar cada ejemplo de entrenamiento de memoria: eso es overfitting, que vimos en «Deep learning a secas», episodio siete. La primera receta se llama L2: a la pérdida le sumamos lambda medios por la norma de theta al cuadrado. Theta son todos los pesos; lambda, un escalar que dice cuánto castigar. La segunda se llama weight decay: en cada paso, multiplicamos los pesos por uno menos eta por lambda, un número apenas menor que uno. Eta es el learning rate. Las dos empujan hacia cero. La diferencia está en dónde se mete cada una: L2 en la pérdida, y por lo tanto en el gradiente; weight decay directo sobre los pesos.

El gradiente del castigo L2

Veamos qué le hace L2 al gradiente, con tres pesos. Theta es un vector de tres números: dos, menos uno, cero coma cinco. Lambda vale cero coma uno. El castigo es lambda medios por la suma de los cuadrados. Le pedimos a PyTorch su gradiente con backward. Y lo comparamos con lambda por theta, calculado a mano. El castigo vale cero coma dos seis dos cinco. El gradiente es cero coma dos, menos cero coma uno, cero coma cero cinco: exactamente lambda por theta. Cada peso recibe un empujón hacia cero proporcional a su propio tamaño.

theta = torch.tensor([2.0, -1.0, 0.5], requires_grad=True)
lam = 0.1                                       # λ: fuerza del L2

penal = lam / 2 * (theta ** 2).sum()            # (λ/2)·‖θ‖²
penal.backward()

print("θ        :", theta.tolist())
print("λ/2·‖θ‖² :", round(penal.item(), 4))
print("gradiente:", [round(x, 4) for x in theta.grad.tolist()])
print("λ·θ      :", [round(lam * x, 4) for x in theta.tolist()])

θ        : [2.0, -1.0, 0.5]
λ/2·‖θ‖² : 0.2625
gradiente: [0.2, -0.1, 0.05]
λ·θ      : [0.2, -0.1, 0.05]

L2, con los números del snippet

$\|\theta\|^2$: la suma de los cuadrados

derivada de un cuadrado: baja el 2, que se cancela con el medio

$g$: el gradiente de la pérdida sola

Lo mismo, término a término. R de theta es el castigo. Las dobles barras son la norma: la norma al cuadrado es la suma de los cuadrados de los pesos. En el dibujo, cuatro, uno y cero coma veinticinco suman cinco coma veinticinco; por lambda medios, cero coma cero cinco, da cero coma dos seis dos cinco. Derivamos peso por peso. La derivada de theta al cuadrado es dos theta; el dos se cancela con el medio y queda lambda por theta. En el dibujo, la pendiente de la parábola es más empinada cuanto más lejos del cero. Y la nabla de toda la pérdida es g, el gradiente de la pérdida original, un vector del mismo tamaño que theta, más lambda theta. En el dibujo: L2 entra como un sumando más del gradiente. Ese detalle decide todo el episodio.

Parte 2 · En SGD son lo mismo

Si el optimizador solo multiplica el gradiente por η, da igual dónde se ponga el castigo.

Tres caminos, el mismo resultado

Una pérdida de juguete con dos pesos, que quiere llevarlos a tres y menos dos. Learning rate cero coma uno, lambda cero coma cinco. El camino A es L2: sumamos lambda por a al gradiente y damos el paso de SGD. El camino B es weight decay: primero multiplicamos por uno menos eta lambda, después restamos eta por el gradiente de la pérdida sola. El camino C es torch punto optim punto SGD con su parámetro weight decay. Diez pasos, y medimos la diferencia. Los tres terminan en uno coma ocho cero tres uno tres y menos cero coma ocho siete tres nueve seis. El factor es cero coma noventa y cinco, y la diferencia entre A y B, una diezmillonésima: redondeo.

eta, lam = 0.1, 0.5
a = torch.tensor([1.0, 1.0])                    # A: L2 dentro del gradiente
b = torch.tensor([1.0, 1.0])                    # B: achicar aparte
c = torch.tensor([1.0, 1.0], requires_grad=True)
opt = torch.optim.SGD([c], lr=eta, weight_decay=lam)
for t in range(10):
    a = a - eta * (grad_perdida(a) + lam * a)
    b = (1 - eta * lam) * b - eta * grad_perdida(b)
    c.grad = grad_perdida(c.detach()); opt.step()
print("A, L2      :", a.numpy().round(5))
print("B, decay   :", b.numpy().round(5))
print("C, torch   :", c.detach().numpy().round(5))
print(f"factor 1 − ηλ = {1 - eta * lam}")
print(f"máx |A − B| = {(a - b).abs().max():.1e}")

A, L2      : [ 1.80313 -0.87396]
B, decay   : [ 1.80313 -0.87396]
C, torch   : [ 1.80313 -0.87396]
factor 1 − ηλ = 0.95
máx |A − B| = 1.2e-07

La misma cuenta, reordenada

SGD con L2: el paso de siempre, con el gradiente nuevo

distribuimos el $\eta$

sacamos factor común $\theta$: eso es weight decay

El paso de SGD con L2: theta menos eta por el gradiente, que ahora es g más lambda theta. En el dibujo, theta y los dos términos que se restan, cada uno por menos eta, menos cero coma uno. Distribuimos el eta: queda theta, menos eta g, menos eta lambda theta. En el dibujo, los tres coeficientes: theta aparece dos veces, con uno y con menos cero coma cero cinco. Sacamos factor común theta: uno menos eta lambda, por theta, menos eta g. Con estos números, cero coma noventa y cinco. Eso es literalmente weight decay: achicar los pesos y después dar el paso de la pérdida sola. En el dibujo, A, B y C del snippet, en el mismo número.

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

SGD: iguales · Adam: no · «L2 regularization and weight decay regularization are equivalent for standard stochastic gradient descent (when rescaled by the learning rate), but as we demonstrate this is not the case for adaptive gradient algorithms, such as Adam.»

Esto es la primera frase importante del paper de AdamW, de Loshchilov y Hutter, de dos mil diecisiete. Dicen que L2 y weight decay son equivalentes para el stochastic gradient descent estándar, cuando se reescalan por el learning rate, que es el eta por lambda que acabamos de ver. Pero que eso no vale para los algoritmos de gradiente adaptativos, como Adam.

Parte 3 · En Adam se separan

Adam divide cada coordenada del gradiente por su propia escala. El término de L2 también pasa por esa división.

Dos pesos, dos escalas de gradiente

En Álgebra lineal para LLMs, episodio diez, vimos el caso extremo: gradiente cero, y Adam con L2 achicaba diez veces más que AdamW. Ahora, algo más realista. Los valores del episodio cinco: beta uno cero coma nueve, beta dos cero coma nueve nueve nueve, épsilon diez a la menos ocho; learning rate cero coma cero uno, lambda cero coma uno. Dos pesos que arrancan en uno. La pérdida les da un gradiente que cambia de signo en cada paso, con promedio cero: diez para A, cero coma cero uno para B. La versión AdamW achica aparte. La versión con L2 suma lambda theta al gradiente. Y el resto es Adam, igual en los dos. Trescientos pasos. Con L2, A sigue en cero coma nueve ocho ocho, y B llegó a cero. Con AdamW, los dos en cero coma siete cinco cuatro.

eta, lam, b1, b2, eps = 0.01, 0.1, 0.9, 0.999, 1e-8
def adam(decay_aparte, pasos=300):
    th, m, v = torch.tensor([1.0, 1.0]), torch.zeros(2), torch.zeros(2)
    hist = [th.clone()]
    for t in range(1, pasos + 1):
        g = ruido(t)                            # A: escala 10; B: escala 0,01
        if decay_aparte: th = th * (1 - eta * lam)   # AdamW
        else:            g = g + lam * th            # Adam + L2
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g ** 2
        mh, vh = m / (1 - b1 ** t), v / (1 - b2 ** t)
        th = th - eta * mh / (vh.sqrt() + eps)
        hist.append(th.clone())
    return torch.stack(hist)
l2, w = adam(False), adam(True)
for nombre, h in [("Adam + L2", l2), ("AdamW    ", w)]:
    print(nombre, "paso 300 → A:", round(h[-1, 0].item(), 3),
          " B:", round(h[-1, 1].item(), 3))

Adam + L2 paso 300 → A: 0.988  B: 0.0
AdamW     paso 300 → A: 0.754  B: 0.754

Por qué: el castigo pasa por √v

el castigo es igual para los dos; el gradiente, no

$\sqrt{\hat v}$: la escala típica de $g'$

el decay no pasa por la división

El paso de Adam del episodio cinco: eta por m sombrero, el promedio del gradiente, dividido la raíz de v sombrero, el del gradiente al cuadrado, más épsilon. Con L2, el gradiente que entra a Adam es g prima: g más lambda theta. En el dibujo, el castigo vale cero coma uno para los dos; el gradiente, diez y cero coma cero uno. Adam divide por la raíz de v, que es la escala típica de g prima. En A, esa escala es diez: el castigo queda en un uno por ciento de paso. En B, el castigo domina la escala, cero coma uno, y queda un paso entero. En el dibujo, cien veces más para B. AdamW saca el decay de la fracción. En el dibujo, la salida real: A y B iguales.

θ_A y θ_B, paso a paso

A, con L2: casi una recta horizontal. Su gradiente es tan grande que el castigo se pierde en la división. B, con L2: cae hasta cero en unos doscientos pasos. Todo lo que Adam ve es el castigo, normalizado a un paso completo. Con AdamW, A y B hacen exactamente el mismo camino: Adam sin decay es invariante a la escala del gradiente, así que el ruido los mueve igual, y el decay los achica igual, un factor cero coma nueve nueve nueve por paso.

Lo mismo en torch.optim

¿Y PyTorch? Las dos clases tienen un parámetro que se llama igual, weight decay, y eso confunde. Mismo learning rate y mismo lambda para las dos. Les damos a mano el mismo gradiente con ruido. Primero, torch punto optim punto Adam. Después, AdamW. Adam con weight decay da cero coma nueve ocho ocho y cero: es nuestra versión con L2, número por número. AdamW da cero coma siete cinco cuatro y cero coma siete cinco cuatro: es la versión desacoplada.

def correr(Clase):
    th = torch.tensor([1.0, 1.0], requires_grad=True)
    opt = Clase([th], lr=0.01, weight_decay=0.1)
    for t in range(1, 301):
        th.grad = ruido(t)                     # solo el gradiente de la pérdida
        opt.step()
    return th.detach().numpy().round(3)

print("torch.optim.Adam(weight_decay=0.1): ", correr(torch.optim.Adam))
print("torch.optim.AdamW(weight_decay=0.1):", correr(torch.optim.AdamW))

torch.optim.Adam(weight_decay=0.1):  [0.988 0.   ]
torch.optim.AdamW(weight_decay=0.1): [0.754 0.754]

Parte 4 · En una red de verdad

Una red chica que lee dígitos escritos a mano. Mismo λ, los dos optimizadores.

Dígitos: qué queda de cada peso

Los dígitos de ocho por ocho que vienen con scikit-learn, y una red de dos capas lineales. W es la matriz de la primera capa: treinta y dos filas por sesenta y cuatro columnas, una columna por píxel. Mil pasos con learning rate diez a la menos tres y lambda cero coma uno, primero con Adam, después con AdamW. Al final medimos qué fracción de su norma inicial conserva cada columna. Y separamos los píxeles muertos, que valen cero en todas las imágenes: sus pesos nunca reciben gradiente de la pérdida. Hay tres. Con Adam y L2, sus pesos quedaron en cero; los vivos, en cero coma seis siete. Con AdamW, los muertos conservan cero coma noventa, que es cero coma nueve nueve nueve nueve a la mil; los vivos crecieron a tres coma cero cuatro.

def entrenar(Clase, pasos=1000):
    torch.manual_seed(0)
    red = torch.nn.Sequential(torch.nn.Linear(64, 32), torch.nn.ReLU(),
                              torch.nn.Linear(32, 10))
    W = red[0].weight                          # 32 × 64: una columna por píxel
    w0 = W.detach().clone()
    opt = Clase(red.parameters(), lr=1e-3, weight_decay=0.1)
    for t in range(pasos):
        b = torch.randint(0, len(X), (64,))
        opt.zero_grad()
        torch.nn.functional.cross_entropy(red(X[b]), y[b]).backward()
        opt.step()
    r = W.detach().norm(dim=0) / w0.norm(dim=0)    # cuánto quedó de cada columna
    return r[muertos].mean(), r[~muertos].mean()
print("píxeles siempre en 0:", int(muertos.sum()), "de 64")
for Clase in [torch.optim.Adam, torch.optim.AdamW]:
    m, v = entrenar(Clase)
    print(f"{Clase.__name__:5}  queda: muertos {m:.2f}  vivos {v:.2f}")

píxeles siempre en 0: 3 de 64
Adam   queda: muertos 0.00  vivos 0.67
AdamW  queda: muertos 0.90  vivos 3.04

Fracción de la norma que queda

Los cuatro números, como barras. Con L2, los píxeles muertos: nada. Los vivos, dos tercios: el castigo les ganó a la pérdida. Con AdamW, los muertos conservan el noventa por ciento: solo el decay, parejo. Y los vivos crecieron, porque la pérdida pidió pesos más grandes y el decay no se lo impidió.

• muertos · Adam + L2: 0,00
• vivos · Adam + L2: 0,67
• muertos · AdamW: 0,90
• vivos · AdamW: 3,04

Qué dice el experimento

• Mismo λ, regularización muy distinta
• Adam + L2: decay según la escala del gradiente
• AdamW: todos se achican por $(1 - \eta\lambda)$
• En torch, Adam(weight_decay) es L2

Leamos esos números. El mismo lambda, cero coma uno, produce dos regularizaciones muy diferentes. No hay forma de elegir un lambda de Adam con L2 que equivalga a uno de AdamW. Con L2, cuánto se achica cada peso depende de cuánto gradiente recibe: los que casi no reciben gradiente se borran, y los que reciben mucho apenas se regularizan. Es al revés de lo que suele querer uno. Con AdamW, el achique es el mismo para todos, independiente del gradiente: el número que escribís es el que se aplica. Y la trampa práctica: en PyTorch, weight decay en Adam significa L2. Si querés el decay desacoplado, es AdamW.

Las dos recetas, lado a lado

La diferencia, en una pantalla. Adam con L2: el castigo entra al gradiente, se mezcla en m y en v, y termina dividido por la raíz de v, distinta para cada peso. Es lo que hace torch punto optim punto Adam. AdamW: el castigo nunca toca m ni v; multiplica los pesos por uno menos eta lambda, el mismo factor para todos. En SGD, las dos columnas serían idénticas.

Resumiendo

• L2 suma $\lambda\theta$ al gradiente
• Weight decay multiplica $\theta$ por $1 - \eta\lambda$
• En SGD: la misma cuenta, reordenada
• En Adam: $\lambda\theta$ se divide por $\sqrt{\hat v}$
• Con el mismo λ: un peso queda en 0,988 y el otro en 0
• Próximo: AdamW, línea por línea

Papers del episodio

• Loshchilov y Hutter, 2017 — Decoupled Weight Decay Regularization