Adam

Optimizadores, episodio cinco: Adam, adaptive moment estimation. En Alineación, episodio dos, AdamW, que es Adam con weight decay desacoplado, pasó en una sola escena: dos promedios, una raíz, una resta. Hoy desarmamos Adam línea por línea. Es la suma de los dos episodios anteriores: la memoria del momentum y la escala por parámetro de RMSProp, root mean square propagation, más una corrección que arregla el arranque.

Parte 1 · Dos promedios

Un promedio del gradiente y un promedio de su cuadrado. Con nombres de estadística: el primer y el segundo momento.

Los dos momentos

primer momento: el promedio de $g$

segundo momento: el promedio de $g^2$

$\hat m$, $\hat v$: corregidos (parte 2)

m es el primer momento: un promedio móvil exponencial del gradiente, con beta uno, una letra griega, igual a cero coma nueve. Es el momentum del episodio tres, pero el gradiente nuevo entra multiplicado por uno menos beta uno: eso lo vuelve un promedio en vez de una suma. En el dibujo, los pesos. v es el segundo momento: el promedio del gradiente al cuadrado, como en RMSProp. Pero con beta dos igual a cero coma nueve nueve nueve. En el dibujo, la memoria de cada uno: uno sobre uno menos beta. m recuerda unos diez pasos; v, unos mil. El paso: eta, el learning rate, que en el paper se llama alfa, por m sombrero sobre la raíz de v sombrero. Los sombreros son la corrección de la parte dos. Todo es coordenada por coordenada: m y v tienen el tamaño de los pesos.

Adam a mano, contra torch.optim.Adam

Las cinco líneas del algoritmo uno del paper, en una clase. El estado es m, v y el contador de pasos, t. El primer momento. El segundo momento, con el gradiente al cuadrado. La corrección: dividir por uno menos beta a la t. Beta a la t es beta multiplicado por sí mismo t veces. Y el paso, con épsilon diez a la menos ocho para no dividir por cero. Al lado, torch punto optim punto Adam con learning rate cero coma uno y los demás valores por defecto, que son los mismos. Cuarenta pasos en el valle. Los dos terminan en menos uno coma cero cinco cero cuatro, menos cero coma cero seis siete. Diferencia máxima: cuatro diezmillonésimas, redondeo.

Cada línea, su término:

• self.m = self.b1 * self.m → $m_t = \beta_1 m_{t-1} + (1 - \beta_1)\, g_t$ — primer momento
• self.v = self.b2 * self.v → $v_t = \beta_2 v_{t-1} + (1 - \beta_2)\, g_t^2$ — segundo momento
• m_hat = self.m / (1 - self.b1 ** self.t) → $\hat m_t = \frac{m_t}{1 - \beta_1^{\,t}},\quad \hat v_t = \frac{v_t}{1 - \beta_2^{\,t}}$ — corrección de sesgo
• return p - self.lr * m_hat → $\theta_t = \theta_{t-1} - \eta\, \frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon}$ — el paso

class Adam:
    def __init__(self, lr, b1=0.9, b2=0.999, eps=1e-8):
        self.lr, self.b1, self.b2, self.eps = lr, b1, b2, eps
        self.m, self.v, self.t = 0.0, 0.0, 0
    def paso(self, p, g):
        self.t += 1
        self.m = self.b1 * self.m + (1 - self.b1) * g        # 1er momento
        self.v = self.b2 * self.v + (1 - self.b2) * g ** 2   # 2do momento
        m_hat = self.m / (1 - self.b1 ** self.t)             # corrección
        v_hat = self.v / (1 - self.b2 ** self.t)
        return p - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

mano, p = Adam(lr=0.1), np.array([-4.5, 1.5]); camino = [p]
w = torch.tensor([-4.5, 1.5], requires_grad=True)
opt = torch.optim.Adam([w], lr=0.1)
for t in range(40):
    p = mano.paso(p, grad(p)); camino.append(p)
    opt.zero_grad(); valle(w).backward(); opt.step()
print("mano: ", p.round(4), " torch:", w.detach().numpy().round(4))
print(f"diferencia máxima: {np.abs(p - w.detach().numpy()).max():.1e}")

mano:  [-1.0504 -0.067 ]  torch: [-1.0504 -0.067 ]
diferencia máxima: 4.5e-07

Parte 2 · El sesgo hacia cero

m y v arrancan en cero. En los primeros pasos, los dos promedios dan de menos.

Los promedios sin corregir

La pendiente constante del episodio cuatro: gradiente siempre cero coma cinco, con las betas de Adam. El promedio verdadero de g es cero coma cinco, y el de g al cuadrado, cero coma veinticinco. El paso crudo divide m por la raíz de v, sin corregir. El corregido usa los sombreros. En el paso uno, m vale cero coma cero cinco: diez veces menos de lo que debería. v vale dos coma cinco diezmilésimas: mil veces menos. El paso crudo da tres coma uno seis; el corregido, uno exacto. En el paso diez, el crudo sube a seis coma cinco tres. En el cien, m ya llegó a cero coma cinco pero v sigue muy abajo: tres coma dos cuatro. En el mil, uno coma dos seis. El corregido da uno en todos los pasos.

g, b1, b2 = 0.5, 0.9, 0.999          # gradiente constante, betas de Adam
m = v = 0.0
for t in range(1, 1001):
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g ** 2
    if t in (1, 10, 100, 1000):
        crudo = m / np.sqrt(v)
        corregido = (m / (1 - b1 ** t)) / np.sqrt(v / (1 - b2 ** t))
        print(f"t = {t:4}: m = {m:.3f}  v = {v:.5f}  "
              f"paso crudo {crudo:5.2f}  corregido {corregido:.2f}")

t =    1: m = 0.050  v = 0.00025  paso crudo  3.16  corregido 1.00
t =   10: m = 0.326  v = 0.00249  paso crudo  6.53  corregido 1.00
t =  100: m = 0.500  v = 0.02380  paso crudo  3.24  corregido 1.00
t = 1000: m = 0.500  v = 0.15808  paso crudo  1.26  corregido 1.00

De dónde sale 1 − βᵗ

m desarmado, como el momentum

suma geométrica

los sombreros devuelven el promedio verdadero

Desarmamos m como el momentum: uno menos beta uno, por la suma de los gradientes con pesos beta a la t menos k. Pero la suma tiene solo t términos. En el dibujo, en el paso tres solo hay tres gradientes, y sus pesos suman cero coma dos siete uno, no uno: lo que falta es el cero inicial. Con gradiente constante, g sale afuera. Queda una suma geométrica: uno menos beta, por uno más beta más beta al cuadrado, hasta beta a la t menos uno. Al multiplicar, los términos del medio se cancelan de a pares y queda uno menos beta a la t. En el dibujo, cero coma uno, cero coma seis cinco uno y uno: los números del snippet. Así que dividir por uno menos beta a la t devuelve exactamente g. Lo mismo para v, con beta dos; como está más cerca de uno, v tarda más en llenarse y el paso crudo se agranda.

El paso, con y sin corrección

Los primeros cien pasos con gradiente constante. Sin corregir, el paso arranca en tres coma uno seis, trepa hasta más de seis en los primeros diez pasos, y recién después baja, muy despacio. Es la joroba del RMSProp del episodio cuatro, pero más alta y más larga, porque beta dos es cero coma nueve nueve nueve. Un paso seis veces más largo con los pesos recién inicializados puede romper un entrenamiento. Corregido, el paso es uno desde el primer momento: el largo que pediste con eta.

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

α = 0,001 · β₁ = 0,9 · β₂ = 0,999 · ε = 10⁻⁸ · «moment estimates that are biased towards zero, especially during the initial timesteps»

Adam es de Diederik Kingma y Jimmy Ba, subido a arXiv en dos mil catorce y publicado en la conferencia ICLR de dos mil quince. En el algoritmo uno del paper están los valores por defecto que usa PyTorch: alfa, nuestro eta, cero coma cero cero uno, beta uno cero coma nueve, beta dos cero coma nueve nueve nueve y épsilon diez a la menos ocho. Y el problema de esta parte, textual: estimaciones de los momentos sesgadas hacia cero, sobre todo en los primeros pasos y con betas cerca de uno. Lo corrigen con los dos cocientes que vimos.

Parte 3 · Qué hace el paso

El cociente entre los dos momentos dice más que el tamaño del gradiente: dice cuánto se puede confiar en su dirección.

RMSProp (ep. 4) contra Adam

El valle, cuarenta pasos, mismo learning rate. RMSProp, del episodio cuatro: sale en diagonal, con los primeros pasos agrandados por la joroba, y avanza por el fondo hasta menos cero coma seis ocho. Adam: sale con el largo justo, también en diagonal, y la memoria de m lo hace pasarse del fondo y volver en una curva suave. Termina en menos uno coma cero cinco: un poco más atrás, porque no tuvo los pasos inflados del principio. En este valle, la joroba ayudó por casualidad; en una red recién inicializada, suele ser lo contrario.

El largo del paso según el gradiente

Una función que corre Adam sobre una lista de gradientes y devuelve el largo del último paso, con learning rate uno. Primero, gradientes constantes, de cero coma cero cero uno y de mil: seis órdenes de magnitud de diferencia. Después, una señal de uno con ruido de desvío tres, como el gradiente de un minibatch del episodio dos. Y solo ruido, de promedio cero. Doscientos pasos cada uno. Los dos constantes dan paso uno: Adam no mira el tamaño del gradiente. Con señal y ruido, cero coma dos uno cinco. Con solo ruido, cero coma cero cero cinco: casi no se mueve.

def paso_adam(gs, b1=0.9, b2=0.999):   # largo del paso, con lr = 1
    m = v = 0.0
    for t, g in enumerate(gs, start=1):
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g ** 2
    return abs(m / (1 - b1 ** t)) / np.sqrt(v / (1 - b2 ** t))

n = 200
for nombre, gs in [("constante 0,001", np.full(n, 0.001)),
                   ("constante 1000 ", np.full(n, 1000.0)),
                   ("señal + ruido  ", 1 + 3 * rng.normal(size=n)),
                   ("solo ruido     ", 3 * rng.normal(size=n))]:
    print(f"{nombre}: paso {paso_adam(gs):.3f}")

constante 0,001: paso 1.000
constante 1000 : paso 1.000
señal + ruido  : paso 0.215
solo ruido     : paso 0.005

El paso como cociente señal sobre ruido

promedios de los últimos pasos

el promedio del cuadrado: media al cuadrado más varianza

$\mu = 1$, $\sigma = 3$: medido 0,215

m sombrero es más o menos el promedio del gradiente en los últimos pasos, y v sombrero el promedio de su cuadrado. En el dibujo, los cuatro casos del snippet. Una identidad de Álgebra, episodio ocho: el promedio del cuadrado es el cuadrado del promedio más la varianza, sigma al cuadrado, que mide cuánto tiembla g alrededor de su promedio, mu. Reemplazando, el paso es mu sobre la raíz de mu al cuadrado más sigma al cuadrado. Sin ruido, da uno: paso completo. Con solo ruido, mu es cero y el paso tiende a cero. Con mu uno y sigma tres, uno sobre raíz de diez, cero coma tres dos. Medimos cero coma dos uno cinco: m sombrero también es un promedio con ruido, así que es una aproximación. Adam avanza decidido cuando el gradiente es consistente y frena cuando es ruido.

Parte 4 · Lo que no resuelve, y lo que cuesta

Adam es el default casi universal. No gana siempre, y paga en memoria.

Rosenbrock: momentum, RMSProp y Adam

La banana del episodio tres, quinientos pasos. El momentum de ese episodio, con learning rate cero coma cero cero uno. Una sola rama para RMSProp y Adam: Adam corrige el sesgo y usa m; RMSProp usa el gradiente crudo y beta dos cero coma nueve. Para cada uno, el mejor learning rate de una grilla corta: cero coma cero dos para RMSProp y cero coma uno para Adam. Momentum queda a cero coma cero nueve cuatro del mínimo. RMSProp, a cero coma cinco cuatro cinco. Adam, a cero coma dos seis cinco. En esta función, momentum gana: el valle curvo premia acumular velocidad a lo largo del fondo, y los métodos adaptativos dan pasos de largo casi fijo.

def grad(p):                                   # Rosenbrock
    x, y = p
    return np.array([-2 * (1 - x) - 400 * x * (y - x * x), 200 * (y - x * x)])

def correr(regla, lr, pasos=500, b1=0.9, b2=0.999):
    p, m, v = np.array([-1.2, 1.0]), 0.0, 0.0; camino = [p]
    for t in range(1, pasos + 1):
        g = grad(p)
        if regla == "momentum":
            m = 0.9 * m + g; p = p - lr * m
        else:                                   # RMSProp (b1 = 0) o Adam
            m = b1 * m + (1 - b1) * g; v = b2 * v + (1 - b2) * g ** 2
            mh, vh = (m / (1 - b1 ** t), v / (1 - b2 ** t)) if regla == "Adam" else (g, v)
            p = p - lr * mh / (np.sqrt(vh) + 1e-8)
        camino.append(p)
    return np.array(camino)

caminos = {"momentum": correr("momentum", 0.001), "RMSProp": correr("RMSProp", 0.02, b2=0.9),
           "Adam": correr("Adam", 0.1)}
for nombre, c in caminos.items():
    print(f"{nombre:8}: a {np.linalg.norm(c[-1] - [1, 1]):.3f} de (1, 1)")

momentum: a 0.094 de (1, 1)
RMSProp : a 0.545 de (1, 1)
Adam    : a 0.265 de (1, 1)

Rosenbrock: quinientos pasos

Un punto cada diez pasos. Momentum, el del episodio tres: toma velocidad por el fondo y llega cerca de la estrella. RMSProp avanza por la curva a paso parejo y se queda a mitad de camino. Adam va por el mismo fondo y llega más lejos que RMSProp, pero no tanto como momentum. Ningún optimizador gana en todas las superficies. Adam se volvió el default por otra razón: con sus valores por defecto funciona razonablemente en problemas muy distintos, sin buscar el learning rate de cero.

Cuánta memoria agrega cada uno

Una red chica de dos capas lineales: de sesenta y cuatro a doscientos cincuenta y seis, y de ahí a diez. Para cada optimizador, un paso, que crea su estado interno. Juntamos los tensores que guarda por parámetro. Y contamos cuántos números son. La red tiene diecinueve mil doscientos diez pesos. SGD, stochastic gradient descent, no guarda nada. Momentum guarda la velocidad: un número por peso. Adam guarda m y v: treinta y ocho mil cuatrocientos veinte, el doble de los pesos. Con siete mil millones de parámetros en float32, son cincuenta y seis gigabytes solo de estado. Es el problema que atacan los optimizadores del episodio diez.

modelo = nn.Sequential(nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 10))
n = sum(p.numel() for p in modelo.parameters())
for nombre, opt in [("SGD", torch.optim.SGD(modelo.parameters(), lr=0.1)),
                    ("momentum", torch.optim.SGD(modelo.parameters(), lr=0.1,
                                                 momentum=0.9)),
                    ("Adam", torch.optim.Adam(modelo.parameters()))]:
    modelo(torch.randn(4, 64)).sum().backward()
    opt.step()                                  # crea el estado
    estado = [v for s in opt.state.values() for v in s.values()
              if torch.is_tensor(v) and v.numel() > 1]
    extra = sum(v.numel() for v in estado)
    print(f"{nombre:8}: {extra:6} números extra ({extra / n:.0f}× los pesos)")
print(f"pesos:    {n:6}")

SGD     :      0 números extra (0× los pesos)
momentum:  19210 números extra (1× los pesos)
Adam    :  38420 números extra (2× los pesos)
pesos:     19210

Cinco optimizadores, una tabla

Los cinco episodios en una tabla. SGD no guarda nada, el paso es eta por g. Momentum guarda una velocidad y da el paso con ella. AdaGrad, adaptive gradient, guarda la suma de los cuadrados y divide por su raíz. RMSProp cambia la suma por un promedio móvil. Y Adam guarda los dos promedios, los corrige y los divide. Cada fila agrega una idea a la anterior.

Lo que hay que llevarse

• $m$: promedio de $g$ (momentum)
• $v$: promedio de $g^2$ (RMSProp)
• $1 - \beta^t$ corrige el arranque en cero
• El paso: largo $\eta$ si $g$ es consistente, chico si es ruido

m es un promedio móvil del gradiente: la memoria del momentum, pero normalizada. v es un promedio móvil del gradiente al cuadrado: la escala por parámetro de RMSProp. Los dos arrancan en cero, y dividir por uno menos beta a la t los devuelve a su valor verdadero desde el primer paso. El cociente entre los dos hace que el paso mida eta cuando el gradiente apunta siempre para el mismo lado, y casi cero cuando es puro ruido, sin importar su tamaño.

Resumiendo

• $m \leftarrow \beta_1 m + (1-\beta_1) g$, $v \leftarrow \beta_2 v + (1-\beta_2) g^2$
• Sin corregir: paso 3,16 en $t=1$ y 6,53 en $t=10$
• Corregido: $\hat m = m/(1-\beta_1^t)$, paso 1 siempre
• torch.optim.Adam: igual a mano (4,5e-7)
• Memoria: 2× los pesos; Rosenbrock: momentum gana
• Próximo: weight decay no es L2

Papers del episodio

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