Schedules

Optimizadores, con foco en AdamW. Episodio ocho: schedules. Hasta ahora el learning rate fue un número fijo. En casi todo entrenamiento de verdad cambia con cada paso: arranca chico, sube y después baja. Hoy armamos ese recorrido con torch punto optim punto lr scheduler, lo comparamos con la fórmula a mano y con la librería transformers, y medimos por qué Adam, adaptive moment estimation, necesita arrancar despacio.

Por qué el learning rate cambia

• Al principio, Adam todavía no conoce la escala
• En el medio, pasos grandes para avanzar
• Al final, pasos chicos para asentarse
• Schedule: el learning rate como función del paso $t$

Tres momentos de un entrenamiento, tres necesidades distintas. Al principio, los promedios de Adam se armaron con muy pocos gradientes: todavía no saben la escala de cada peso. Conviene no confiar mucho en ellos. En el medio, queremos el learning rate más grande que no diverja, como en el episodio uno. Al final, el ruido del minibatch, del episodio dos, hace temblar los pesos alrededor del mínimo, y ese temblor es proporcional al learning rate: bajarlo los deja asentarse. Un schedule es justamente eso: una función que da el learning rate eta sub t en cada paso t.

Parte 1 · Warmup lineal y coseno

La receta más común: subir en línea recta durante unos pasos y después bajar con forma de coseno.

Warmup y coseno con torch.optim.lr_scheduler

Mil pasos en total, con AdamW y un learning rate máximo de diez a la menos tres. LinearLR es el warmup: arranca en el uno por ciento del máximo y sube en línea recta durante cien pasos. CosineAnnealingLR es la bajada: novecientos pasos hasta un mínimo de diez a la menos cinco. SequentialLR los encadena: el primero hasta el paso cien, el segundo después. get last lr devuelve el learning rate que va a usar este paso. Y el scheduler avanza siempre después de opt punto step. En el paso cero, diez a la menos cinco. En el cincuenta, la mitad del camino. En el cien, el máximo. En el quinientos cincuenta, otra vez la mitad, y al final, el mínimo. Sumados, el warmup aporta cero coma cero cinco y el coseno, cero coma cuatro cinco cinco.

w = torch.zeros(10, requires_grad=True)
opt = torch.optim.AdamW([w], lr=1e-3)
warmup = LinearLR(opt, start_factor=0.01, end_factor=1.0, total_iters=100)
coseno = CosineAnnealingLR(opt, T_max=900, eta_min=1e-5)
sched = SequentialLR(opt, [warmup, coseno], milestones=[100])
lrs = []
for paso in range(1000):
    lrs.append(sched.get_last_lr()[0])        # el lr que usa este paso
    w.grad = torch.randn(10); opt.step()
    sched.step()                              # después de opt.step()
for t in [0, 50, 100, 550, 999]:
    print(f"paso {t:3}: lr = {lrs[t]:.6f}")
print(f"suma: warmup {sum(lrs[:100]):.4f}, coseno {sum(lrs[100:]):.4f}")

paso   0: lr = 0.000010
paso  50: lr = 0.000505
paso 100: lr = 0.001000
paso 550: lr = 0.000505
paso 999: lr = 0.000010
suma: warmup 0.0500, coseno 0.4550

El learning rate en cada paso

La salida real del scheduler, paso por paso, en milésimas. Una rampa corta: en cien pasos sube del uno por ciento al máximo. Ese es el warmup. Después, el coseno: baja despacio al principio, más rápido en el medio, y vuelve a frenar cerca del final, como la mitad de una ola. El área bajo la curva es cuánto se mueven los pesos en total: menos de la mitad que con el máximo fijo durante mil pasos.

Las dos fórmulas, con los números del snippet

warmup lineal: $f_0 = 0{,}01$, $T_w = 100$ pasos

coseno: $t$ cuenta desde el fin del warmup; $T = 900$

justo la mitad entre el máximo y el mínimo

El warmup lineal. Eta máximo es el learning rate configurado; f sub cero, el uno por ciento inicial; T sub w, los cien pasos de warmup. En el paso cincuenta, cero coma cero uno más la mitad de cero coma noventa y nueve, por el máximo: cero coma cero cero cero cinco cero cinco. En el dibujo, la recta en tres puntos. El coseno. Coseno de pi por t sobre T va de uno a menos uno; uno más eso, sobre dos, va de uno a cero. En el dibujo, ese factor en cinco puntos: uno, cero coma ocho cinco, cero coma cinco, cero coma uno cinco, cero. Multiplica la distancia entre el máximo y el mínimo. En la mitad, el coseno de pi medios es cero, y el learning rate queda justo a mitad de camino. En el dibujo, la salida del snippet.

El coseno, a mano contra torch

Comprobemos que CosineAnnealingLR es esa fórmula y nada más. La función a mano, con los mismos tres números. Coseno de pi por t sobre T, con math punto cos. El scheduler de PyTorch, sobre un optimizador cualquiera. Y en cada uno de los novecientos un pasos, la diferencia en valor absoluto. En t igual a cero, diez a la menos tres; en cuatrocientos cincuenta, cero coma cero cero cero cinco cero cinco; en novecientos, diez a la menos cinco. La diferencia máxima es uno por diez a la menos dieciocho: son el mismo número.

eta_max, eta_min, T = 1e-3, 1e-5, 900
def coseno(t):                                # la fórmula, a mano
    return eta_min + 0.5 * (eta_max - eta_min) * (1 + math.cos(math.pi * t / T))

opt = torch.optim.SGD([torch.zeros(1, requires_grad=True)], lr=eta_max)
sched = CosineAnnealingLR(opt, T_max=T, eta_min=eta_min)
dif = 0.0
for t in range(T + 1):
    dif = max(dif, abs(sched.get_last_lr()[0] - coseno(t)))
    opt.step(); sched.step()
print(f"t = 0: {coseno(0):.6f}  t = 450: {coseno(450):.6f}  t = 900: {coseno(900):.6f}")
print(f"máx |torch − fórmula| = {dif:.1e}")

t = 0: 0.001000  t = 450: 0.000505  t = 900: 0.000010
máx |torch − fórmula| = 1.1e-18

Parte 2 · Por qué Adam necesita warmup

Al principio, v̂ es el promedio de muy pocos gradientes al cuadrado. Y dividir por algo mal estimado es peligroso.

El problema de los primeros pasos

el primer paso: entero, en cada peso (episodio siete)

el multiplicador sale de una sola muestra

con pocos pasos, el promedio es de pocas muestras

En el episodio siete vimos que en el primer paso m sombrero es g y v sombrero es g al cuadrado, así que el paso es eta por el signo de g. En el dibujo, gradientes de cero coma cero uno, cero coma cinco y menos cuatro dan pasos idénticos. Adam se mueve con todo, sin saber nada todavía. El multiplicador que Adam le aplica al gradiente es uno sobre la raíz de v sombrero, y en el paso uno, eso es uno sobre el valor absoluto de una sola muestra. En el dibujo: si por azar el gradiente salió chico, el multiplicador es enorme. Con beta dos cero coma nueve nueve nueve, v sombrero termina promediando unos mil pasos; al principio promedia uno, dos, diez. En el dibujo, lo que mide el snippet que sigue.

Cien mil pesos con gradiente de puro ruido

Cien mil pesos, con las betas de siempre. El gradiente es ruido puro, de escala uno para todos: un peso no debería moverse más que otro. Medimos el percentil noventa y nueve del multiplicador: el valor que solo supera el uno por ciento de los pesos. Y el largo promedio del paso, en unidades de eta. En el paso uno, el uno por ciento de los pesos multiplica su gradiente por más de ochenta y uno, y el paso promedio es eta entero, todo sobre ruido. En el paso dos, el percentil baja a diez. En el diez, a dos. En el cien y en el mil, ronda uno, y el paso baja a cero coma uno ocho de eta. El ruido se cancela en m, pero recién después de varios pasos.

n, b1, b2 = 100_000, 0.9, 0.999              # cien mil pesos, gradiente = ruido
m, v = torch.zeros(n), torch.zeros(n)
for t in range(1, 1001):
    g = torch.randn(n)                       # escala real: 1 en todos
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g ** 2
    mh, vh = m / (1 - b1 ** t), v / (1 - b2 ** t)
    if t in (1, 2, 10, 100, 1000):
        p99 = (1 / vh.sqrt()).quantile(0.99).item()   # 1/√v̂: el multiplicador
        paso = (mh / vh.sqrt()).abs().mean().item()  # |paso| / η
        print(f"t = {t:4}: 1/√v̂ p99 = {p99:5.1f}   |paso|/η = {paso:.2f}")

t =    1: 1/√v̂ p99 =  81.5   |paso|/η = 1.00
t =    2: 1/√v̂ p99 =  10.1   |paso|/η = 0.64
t =   10: 1/√v̂ p99 =   2.0   |paso|/η = 0.27
t =  100: 1/√v̂ p99 =   1.2   |paso|/η = 0.18
t = 1000: 1/√v̂ p99 =   1.1   |paso|/η = 0.18

Percentil 99 del multiplicador 1/√v̂

Las cinco mediciones como barras. En el paso uno, ochenta y uno. En el dos, diez. En el diez, dos. En el cien, uno coma dos. Y en el mil, uno coma uno: ya casi el valor verdadero, uno. Las primeras decenas de pasos son las peligrosas, y es justo donde el warmup pone el learning rate en una fracción del máximo. Una forma de pensarlo: el warmup no deja que esos multiplicadores mal estimados muevan mucho los pesos.

• t = 1: 81,5
• t = 2: 10,1
• t = 10: 2,0
• t = 100: 1,2
• t = 1000: 1,1

Liu et al., 2019 — On the Variance of the Adaptive Learning Rate and Beyond · arXiv:1908.03265

varianza alta al principio · «we identify a problem of the adaptive learning rate (i.e., it has problematically large variance in the early stage)» … «warmup works as a variance reduction technique»

Esta lectura del warmup es la de Liu y coautores, en dos mil diecinueve, el paper que presentó RAdam. Identifican un problema del learning rate adaptativo: tiene una varianza problemáticamente grande en la etapa temprana. Y sugieren que el warmup funciona como una técnica para reducir esa varianza. Es lo que acabamos de medir con el percentil noventa y nueve.

RAdam: el warmup adentro del optimizador

RAdam, rectified Adam, está en torch punto optim. El mismo gradiente en todos los pasos, cero coma cinco y menos cuatro. Adam contra RAdam, learning rate cero coma uno. Quinientos pasos, guardando cuánto se mueve cada peso. Y miramos algunos. Adam da siempre un paso de cero coma uno en cada coordenada, desde el primero. RAdam, en los primeros cinco pasos, da menos cero coma cero cinco y cero coma cuatro: eta por el gradiente, sin dividir por la raíz de v. Es SGD, stochastic gradient descent, con momentum. En el sexto empieza a dividir, pero con un factor que arranca chico: tres milésimas. En el cincuenta, quince milésimas. En el quinientos, cuarenta y ocho, todavía la mitad. Un warmup automático.

g = torch.tensor([0.5, -4.0])                 # el mismo gradiente en cada paso
for Clase in [torch.optim.Adam, torch.optim.RAdam]:
    w = torch.zeros(2, requires_grad=True)
    opt = Clase([w], lr=0.1)
    pasos = []
    for t in range(1, 501):
        antes = w.detach().clone()
        w.grad = g.clone(); opt.step()
        pasos.append((w.detach() - antes).tolist())
    print(Clase.__name__)
    for t in (1, 5, 6, 50, 500):
        print(f"  paso {t}: {[round(x, 3) for x in pasos[t - 1]]}")

Adam
  paso 1: [-0.1, 0.1]
  paso 5: [-0.1, 0.1]
  paso 6: [-0.1, 0.1]
  paso 50: [-0.1, 0.1]
  paso 500: [-0.1, 0.1]
RAdam
  paso 1: [-0.05, 0.4]
  paso 5: [-0.05, 0.4]
  paso 6: [-0.003, 0.003]
  paso 50: [-0.015, 0.015]
  paso 500: [-0.048, 0.048]

Dos formas de arrancar despacio

Las dos soluciones, lado a lado. El warmup vive afuera del optimizador: es el learning rate el que arranca chico, y sirve igual para Adam, AdamW o cualquier otro. A cambio, hay que elegir cuántos pasos dura. RAdam lo resuelve adentro: no divide por la raíz de v hasta tener unos pocos gradientes, y después aplica un factor que crece solo. Las dos ideas atacan lo mismo: no confiar en v sombrero mientras es el promedio de pocas muestras.

Parte 3 · En la práctica

Lo que vas a encontrar en código de fine-tuning, y lo que hay que cuidar.

El schedule de la librería transformers

En código de fine-tuning con Hugging Face vas a ver esto. get cosine schedule with warmup, de la librería transformers. Cien pasos de warmup. Mil pasos en total. Por dentro devuelve un LambdaLR de PyTorch: un factor por paso que multiplica el learning rate. Lo comparamos con nuestro SequentialLR. Dos diferencias. Transformers arranca en cero, no en el uno por ciento: el paso cero no mueve los pesos. Y termina en cero, no en diez a la menos cinco. En el medio, los mismos números: la mitad en el paso cincuenta, el máximo en el cien.

w = torch.zeros(10, requires_grad=True)
opt = torch.optim.AdamW([w], lr=1e-3)
sched = get_cosine_schedule_with_warmup(opt, num_warmup_steps=100,
                                        num_training_steps=1000)
hf = []
for paso in range(1000):
    hf.append(sched.get_last_lr()[0])
    w.grad = torch.randn(10); opt.step(); sched.step()
print("paso   transformers   SequentialLR")
for t in [0, 1, 50, 100, 550, 999]:
    print(f"{t:4}   {hf[t]:.6f}       {nuestro[t]:.6f}")

paso   transformers   SequentialLR
   0   0.000000       0.000010
   1   0.000010       0.000020
  50   0.000500       0.000505
 100   0.001000       0.001000
 550   0.000500       0.000505
 999   0.000000       0.000010

Los primeros 120 pasos, las dos librerías

Una lupa sobre el warmup. Nuestro SequentialLR arranca apenas arriba de cero y sube en recta hasta el paso cien. Transformers, en punteado, arranca exactamente en cero, y desde ahí son casi la misma recta. Después del cien, los dos bajan con el mismo coseno. Para el entrenamiento, la diferencia es despreciable; lo que importa es tener warmup.

Cuánto se mueven los pesos en total

$\Sigma$: la suma de los mil learning rates

un triángulo: la base por la altura media

el coseno tiene la misma área que su valor medio

Un número útil para comparar schedules: la suma de todos los learning rates, el área bajo la curva. La sigma mayúscula es la suma. Con eta fijo en diez a la menos tres durante mil pasos, da uno. En el dibujo, nuestro schedule suma cero coma cinco cero cinco: la mitad. El warmup es un triángulo, casi: cien pasos con un promedio de la mitad del máximo, cero coma cero cinco. En el dibujo, apenas el diez por ciento del recorrido. Y el coseno: como la media ola es simétrica, su promedio es el punto medio entre el máximo y el mínimo. Novecientos pasos por cero coma cero cero cero cinco cero cinco da cero coma cuatro cinco cinco. En el dibujo, la cuenta y la suma del snippet.

Detalles que rompen entrenamientos

• sched.step() después de opt.step()
• Contar en pasos, no en epochs
• Guardar el scheduler en el checkpoint
• Cambiar el schedule cambia el weight decay de AdamW

Cuatro detalles. El orden: primero opt punto step, después sched punto step. Al revés, PyTorch avisa y el primer learning rate se saltea. Los schedules de fine-tuning se cuentan en pasos de optimizador, no en epochs: si usás gradient accumulation, un paso son varios batches. Si guardás un checkpoint para seguir después, guardá también el state dict del scheduler, o al retomar vuelve al warmup. Y del episodio siete: en PyTorch el decay de AdamW es learning rate por weight decay, así que el schedule también lo apaga al final y lo achica durante el warmup.

Resumiendo

• Schedule: $\eta_t$ distinto en cada paso
• Warmup lineal: del 1 % al máximo en 100 pasos
• Coseno: $\eta_{\min} + \tfrac12(\eta_{\max} - \eta_{\min})(1 + \cos\tfrac{\pi t}{T})$
• Adam al principio: $1/\sqrt{\hat v}$ de pocas muestras, percentil 99 = 81,5
• RAdam y el warmup reducen esa varianza
• Próximo: gradient clipping, épsilon y precisión mixta

Papers del episodio

• Liu et al., 2019 — On the Variance of the Adaptive Learning Rate and Beyond