Momentum

Optimizadores, episodio tres: momentum. En el episodio uno, el valle alargado nos dejó un problema: el learning rate que no explota en la dirección empinada es lentísimo en la plana. Hoy, la primera solución: que el optimizador recuerde los gradientes anteriores. Y una variante que mira un paso adelante, la de Nesterov.

Parte 1 · Gradientes que se repiten

Antes de cambiar la regla, miremos los gradientes que produce descenso en el valle.

Los gradientes de descenso, en el valle

El valle del episodio uno, cero coma uno x al cuadrado más tres y al cuadrado, con el mismo punto de partida y learning rate cero coma tres. En cada paso imprimimos el gradiente, sus dos componentes. Y damos el paso de siempre. Mirá las dos columnas. La de x: menos cero coma nueve, menos cero coma ocho cuatro seis, menos cero coma siete nueve cinco. Siempre negativa y casi igual: apunta para el mismo lado paso tras paso. La de y: nueve, menos siete coma dos, cinco coma siete seis, menos cuatro coma seis. Cambia de signo cada vez: es el rebote entre las paredes.

p, lr = np.array([-4.5, 1.5]), 0.3        # el valle del episodio uno
grads = []
for t in range(6):
    g = np.array([0.2, 6.0]) * p          # gradiente: (0,2x, 6y)
    grads.append(g)
    print(f"paso {t}: g = ({g[0]:6.3f}, {g[1]:6.3f})")
    p = p - lr * g                        # descenso por gradiente

paso 0: g = (-0.900,  9.000)
paso 1: g = (-0.846, -7.200)
paso 2: g = (-0.795,  5.760)
paso 3: g = (-0.748, -4.608)
paso 4: g = (-0.703,  3.686)
paso 5: g = (-0.661, -2.949)

Recordar los gradientes

$\mu$: cuánto se recuerda, entre 0 y 1

reemplazamos $v_{t-1}$ por su definición, una y otra vez

mismo signo: se acumulan

signos alternados: se cancelan

Momentum agrega v, la velocidad: un vector del tamaño de los pesos. Cada paso, v es mu por la v anterior más el gradiente nuevo, g sub t, y los pesos, w, se mueven con v. Mu es un escalar entre cero y uno; el clásico, cero coma nueve. Reemplazando v sub t menos uno por su propia definición, hacia atrás, aparece la suma de los gradientes pasados, pesados por mu, mu al cuadrado, mu al cubo. En el dibujo: los viejos se van apagando. Con los cinco gradientes del snippet, en x: todos negativos, se suman. En el dibujo, menos tres coma cero tres: cuatro veces y media el último. En y: los signos se alternan, y la suma da menos tres coma ocho nueve, apenas más que el último, aunque cada gradiente valía entre tres y nueve. El avance se amplifica y el rebote no.

Momentum a mano, contra torch.optim.SGD

Una clase de cinco líneas. Guarda el learning rate, mu y la velocidad, que arranca en cero. El paso actualiza la velocidad: mu por la vieja más el gradiente. Y mueve los pesos con la velocidad. Al lado, PyTorch: el mismo SGD, stochastic gradient descent, del episodio uno, con el argumento momentum en cero coma nueve. Learning rate cero coma cero tres, cuarenta pasos en el valle. Los dos terminan en cero coma cero ocho cuatro uno, cero coma cero cuatro seis uno. PyTorch guarda la velocidad exactamente así: en su documentación la llama b, y la suma no lleva ningún factor delante del gradiente. Distancia al mínimo: cero coma cero nueve seis.

class Momentum:
    def __init__(self, lr, mu):
        self.lr, self.mu, self.v = lr, mu, 0.0
    def paso(self, p, g):
        self.v = self.mu * self.v + g             # recordar los gradientes
        return p - self.lr * self.v               # moverse con la suma

opt_mano = Momentum(lr=0.03, mu=0.9)
p = np.array([-4.5, 1.5]); camino = [p]
w = torch.tensor([-4.5, 1.5], requires_grad=True)
opt = torch.optim.SGD([w], lr=0.03, momentum=0.9)
for t in range(40):
    p = opt_mano.paso(p, grad(p)); camino.append(p)
    opt.zero_grad(); (0.1 * w[0]**2 + 3 * w[1]**2).backward(); opt.step()
print("mano: ", p.round(4), " torch:", w.detach().numpy().round(4))
print(f"distancia al mínimo: {np.linalg.norm(p):.3f}")

mano:  [0.0841 0.0461]  torch: [0.0841 0.0461]
distancia al mínimo: 0.096

Episodio uno contra momentum

Cuarenta pasos cada uno. Si viste Deep learning a secas, episodio siete, este dibujo te suena: ahí pasó en una escena, y acá lo desarmamos. Descenso, con el learning rate cero coma tres del episodio uno: rebota entre las paredes y avanza despacio por el fondo; queda a casi cuatro décimas. Momentum, con un learning rate diez veces menor: la bolita baja en curva, se pasa un poco, vuelve, y la velocidad acumulada en x la lleva por el fondo hasta la estrella. Termina a menos de una décima.

Parte 2 · Cuánto gana

Para pérdidas cuadráticas como el valle, la ganancia se puede calcular exactamente.

La raíz del condition number

el condition number del episodio uno

lo que queda del error después de cada paso

$\kappa$ pasa a $\sqrt{\kappa}$

Kappa, la letra griega, es el condition number: la curvatura más grande, lambda máximo, dividida la más chica. En el valle, seis sobre cero coma dos, treinta. Con el mejor learning rate posible, descenso achica el error a kappa menos uno sobre kappa más uno por paso: veintinueve treinta y unavos, cero coma nueve tres cinco. Solo saca un seis y medio por ciento por paso. En el dibujo, la barra de descenso. Con momentum, eligiendo bien eta y mu, la fórmula es la misma pero con la raíz de kappa: raíz de treinta es cinco coma cuatro ocho, y queda cero coma seis nueve uno. En el dibujo de abajo, lo que significa en pasos: ciento veintisiete contra treinta y uno.

Los parámetros óptimos, y la prueba

Para descenso, el mejor learning rate es dos sobre la suma de las curvaturas: equilibra la dirección plana y la empinada. Para momentum, Goh da el learning rate óptimo con las raíces de las curvaturas. Y el mu óptimo. Después, corremos los dos y contamos los pasos hasta quedar a una milésima del mínimo. Descenso: learning rate cero coma tres dos tres, ritmo cero coma nueve tres cinco. Momentum: learning rate cero coma cuatro siete siete, mu cero coma cuatro siete ocho, ritmo cero coma seis nueve uno. Y los pasos: ciento veintisiete contra treinta y uno. Cuatro veces menos.

l1, ln = 0.2, 6.0                         # curvaturas del valle
k = ln / l1                               # condition number: 30
lr_gd = 2 / (l1 + ln)                     # el mejor lr para descenso
lr_m = (2 / (np.sqrt(l1) + np.sqrt(ln))) ** 2      # Goh, 2017
mu_m = ((np.sqrt(ln) - np.sqrt(l1)) / (np.sqrt(ln) + np.sqrt(l1))) ** 2
print(f"descenso: lr {lr_gd:.3f}  ritmo {(k - 1) / (k + 1):.3f}")
print(f"momentum: lr {lr_m:.3f}  mu {mu_m:.3f}  ritmo "
      f"{(np.sqrt(k) - 1) / (np.sqrt(k) + 1):.3f}")

def pasos_hasta(lr, mu, tol=1e-3):
    p, v, t = np.array([-4.5, 1.5]), 0.0, 0
    while np.linalg.norm(p) > tol:
        v = mu * v + grad(p); p = p - lr * v; t += 1
    return t
print("pasos hasta 0,001:", pasos_hasta(lr_gd, 0.0), "contra", pasos_hasta(lr_m, mu_m))

descenso: lr 0.323  ritmo 0.935
momentum: lr 0.477  mu 0.478  ritmo 0.691
pasos hasta 0,001: 127 contra 31

El mejor descenso contra el mejor momentum

Los dos con sus parámetros óptimos, cuarenta pasos. Descenso con cero coma tres dos tres: sigue rebotando y a los cuarenta pasos está a más de tres décimas. Momentum con los valores de Goh: los primeros saltos son grandes, de pared a pared, pero se amortiguan rápido y la bolita llega a la estrella. Con momentum, el learning rate óptimo es más grande, y la memoria frena el rebote que ese learning rate provocaría solo.

Goh, 2017 — Why Momentum Really Works · fuente

κ → √κ · «Momentum is a heavy ball rolling down the same hill» · «we have essentially square rooted the condition number!»

El artículo de Gabriel Goh que abrimos en el episodio uno. Su imagen para momentum es una bola pesada que rueda por la misma colina: el gradiente la empuja, y lo que ya venía moviéndose sigue. Y su resultado, con signos de exclamación: con un mínimo de esfuerzo extra, le sacamos la raíz cuadrada al condition number. Con kappa treinta es cuatro veces menos pasos; en redes reales, donde kappa es enorme, la diferencia es mucho mayor.

Parte 3 · Nesterov: mirar antes de saltar

Si de todos modos vamos a movernos mu por v, mejor medir el gradiente donde vamos a caer.

Dos formas de escribir lo mismo

momentum clásico (ecs. 1 y 2 del paper)

Nesterov (ecs. 3 y 4)

cambio de variable: la forma de PyTorch

Sutskever y sus coautores lo escriben con otra convención: v ya es el paso, con el signo menos adentro, y épsilon es el learning rate. Es lo mismo que la b de PyTorch, cambiada de signo y multiplicada por el learning rate. Momentum clásico: la velocidad nueva es mu por la vieja menos épsilon por el gradiente en theta, que son los pesos actuales. En el dibujo, se suman. Nesterov cambia un solo detalle: el gradiente se mide en theta más mu v, el punto al que la velocidad te va a llevar igual. PyTorch no guarda theta: guarda ese punto adelantado, fi. Reemplazando theta por fi menos mu v en las dos ecuaciones, y agrupando los términos, queda esta forma: un gradiente en fi, más mu por el buffer b. Es el mismo algoritmo, mirado desde otro punto.

El gradiente adelantado, en un paso real

Un paso real de momentum en el valle, el tercero. La velocidad está agrandada cuatro veces y los gradientes treinta, para que se vean. En celeste, mu v: la velocidad trae a la bolita hacia abajo, cruzando el fondo del valle. En ámbar punteado, el gradiente donde está ahora: todavía empuja un poco más hacia abajo, porque todavía no cruzó. En verde, el gradiente medido en el punto adelantado, del otro lado del fondo: empuja hacia arriba. Nesterov frena antes de pasarse. El momentum clásico se entera un paso tarde.

valle, paso 3 · μv ×4, gradientes ×30 · clásico (0.026, -0.038) · Nesterov (0.026, 0.052)

Nesterov a mano, contra PyTorch

Las ecuaciones del paper, tal cual: la velocidad con el gradiente adelantado. Y los pesos se mueven con la velocidad. Al lado, PyTorch con nesterov igual a true, mismo learning rate y mismo mu. Al final, calculamos theta más mu v. Theta más mu v da cero coma cero uno uno cinco, cero coma cero cero uno seis, y PyTorch da exactamente lo mismo: confirma el cambio de variable. Y la distancia al mínimo después de cuarenta pasos es cero coma cero dos seis, contra cero coma cero nueve seis del momentum clásico: casi cuatro veces más cerca.

Cada línea, su término:

• v = mu * v - lr * grad(th + mu * v) → $v \leftarrow \mu v - \varepsilon \nabla f(\theta + \mu v)$ — ecuación 3
• th = th + v → $\theta \leftarrow \theta + v$ — ecuación 4
• nesterov=True → $\phi \leftarrow \phi - \varepsilon\,(\nabla f(\phi) + \mu b)$ — lo que corre PyTorch
• adelantado = th + mu * v → $\phi = \theta + \mu v$ — el punto que guarda PyTorch

mu, lr = 0.9, 0.03
th, v = np.array([-4.5, 1.5]), np.zeros(2)       # Sutskever et al., 2013
w = torch.tensor([-4.5, 1.5], requires_grad=True)
opt = torch.optim.SGD([w], lr=lr, momentum=mu, nesterov=True)
for t in range(40):
    v = mu * v - lr * grad(th + mu * v)          # gradiente adelantado
    th = th + v
    opt.zero_grad(); valle(w).backward(); opt.step()
    camino.append(th.tolist())
adelantado = th + mu * v
print("Sutskever, θ + μv:", adelantado.round(4))
print("torch, nesterov:  ", w.detach().numpy().round(4))
print(f"distancia al mínimo: {np.linalg.norm(th):.3f}")

Sutskever, θ + μv: [0.0115 0.0016]
torch, nesterov:   [0.0115 0.0016]
distancia al mínimo: 0.026

Sutskever, Martens, Dahl y Hinton, 2013 — On the importance of initialization and momentum in deep learning · fuente

μ hasta 0,999 · «We find that both the initialization and the momentum are crucial»

El paper que puso a Nesterov en el vocabulario del deep learning. Sutskever, Martens, Dahl y Hinton, en la conferencia ICML de dos mil trece. Mostraron que SGD con momentum, bien inicializado y con mu creciendo hacia valores como cero coma nueve nueve cinco o cero coma nueve nueve nueve, entrenaba redes profundas y recurrentes que antes se creía que necesitaban métodos de segundo orden. Su conclusión, textual: tanto la inicialización como el momentum son cruciales. Las ecuaciones de la escena anterior son las tres y cuatro de ese paper.

Momentum contra Nesterov

Mismo learning rate, mismo mu, cuarenta pasos. Momentum clásico: baja, se pasa del fondo, vuelve a pasarse, y se acerca dando curvas. Nesterov: el primer cruce del fondo es más corto, porque frena antes, y llega más derecho. Termina casi cuatro veces más cerca. En el valle, que es una cuadrática, la ventaja es clara.

Un valle curvo: Rosenbrock

Ahora una pérdida que no es cuadrática: la función de Rosenbrock, uno menos x al cuadrado más cien por y menos x al cuadrado, al cuadrado. Tiene el mínimo en uno, uno, al fondo de un valle angosto y curvo con forma de banana. Su gradiente, derivando cada término con la regla de la cadena. Una sola función para los tres métodos: la diferencia es dónde se mide el gradiente. Learning rate cero coma cero cero uno, quinientos pasos. Descenso queda a uno coma seis siete ocho del mínimo. Momentum, a cero coma cero nueve cuatro. Nesterov, a cero coma uno uno cero: acá no le gana al clásico. Su garantía es para funciones convexas, con forma de cuenco; en otras, depende.

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

def correr(mu, nesterov=False, lr=0.001, pasos=500):
    p, v = np.array([-1.2, 1.0]), np.zeros(2); camino = [p]
    for _ in range(pasos):
        g = grad(p + mu * v) if nesterov else grad(p)
        v = mu * v - lr * g; p = p + v; camino.append(p)
    return np.array(camino)

for nombre, c in [("descenso", correr(0.0)), ("momentum", correr(0.9)),
                  ("Nesterov", correr(0.9, nesterov=True))]:
    print(f"{nombre:9}: a {np.linalg.norm(c[-1] - [1, 1]):.3f} de (1, 1)")

descenso : a 1.678 de (1, 1)
momentum : a 0.094 de (1, 1)
Nesterov : a 0.110 de (1, 1)

Rosenbrock: quinientos pasos

Los tres caminos, con un punto cada diez pasos. Descenso baja enseguida al fondo de la banana y después se arrastra por la curva: en quinientos pasos apenas avanzó. Momentum toma velocidad a lo largo del fondo y llega cerca de la estrella. Nesterov hace casi el mismo recorrido. En pérdidas reales, momentum es la mejora grande; Nesterov es un retoque que a veces ayuda y a veces no.

Resumiendo

• Momentum: $v \leftarrow \mu v + g$, el paso usa $v$
• Suma $\mu^k$: lo constante se acumula, el rebote se cancela
• Valle: 127 pasos contra 31 ($\kappa \to \sqrt{\kappa}$)
• Nesterov: el gradiente en $\theta + \mu v$
• torch.optim.SGD con momentum y nesterov: iguales a mano
• Próximo: un learning rate por parámetro

Papers del episodio

• Goh, 2017 — Why Momentum Really Works
• Sutskever, Martens, Dahl y Hinton, 2013 — On the importance of initialization and momentum in deep learning