English · Español

Lab 02 — Competir SGD, Momentum, Adam, AdamW en Rosenbrock¶

🇪🇸 Cuatro optimizadores en la misma pista: la función de Rosenbrock, un valle estrecho y curvo. Animas las trayectorias. Lo que ves es por qué Adam gana en valles mal-condicionados.

Objetivo¶

Implementa cuatro optimizadores como rutinas *_step de función pura en src/minigrad/optim.py. Ejecuta cada uno sobre la función de Rosenbrock f(x, y) = (1 - x)² + 100(y - x²)² desde un inicio común, registra la trayectoria, y produce una gráfica de contornos animada única con las cuatro trayectorias superpuestas.

Planteamiento¶

Fase 03 (álgebra lineal) y los archivos de teoría de la Fase 04.
La función de Rosenbrock: convexa a escala de conjuntos de nivel, pero el mínimo (en (1, 1)) está al final de un valle estrecho y curvo. La función es el test clásico para el comportamiento de optimizadores en regiones mal condicionadas.
matplotlib, numpy.

La función de Rosenbrock¶

def rosenbrock(p):
    x, y = p
    return (1 - x)**2 + 100 * (y - x**2)**2

def rosenbrock_grad(p):
    x, y = p
    dx = -2 * (1 - x) + 100 * 2 * (y - x**2) * (-2 * x)
    dy = 100 * 2 * (y - x**2)
    return np.array([dx, dy])

Mínimo en (1, 1) donde f = 0. Empieza cada optimizador en (-1.5, 2.0).

Tareas¶

Parte A — Implementa los cuatro optimizadores en `src/minigrad/optim.py`¶

Estilo de función pura con diccionarios de estado explícitos.

def sgd_step(params, grads, state, lr):
    """state is unused for vanilla SGD; kept for API uniformity."""
    return params - lr * grads, state

def momentum_step(params, grads, state, lr, beta=0.9):
    v = state.get("v", np.zeros_like(params))
    v = beta * v + grads
    new_params = params - lr * v
    return new_params, {"v": v}

def adam_step(params, grads, state, lr, beta1=0.9, beta2=0.999, eps=1e-8):
    m = state.get("m", np.zeros_like(params))
    v = state.get("v", np.zeros_like(params))
    t = state.get("t", 0) + 1
    m = beta1 * m + (1 - beta1) * grads
    v = beta2 * v + (1 - beta2) * grads**2
    m_hat = m / (1 - beta1**t)
    v_hat = v / (1 - beta2**t)
    new_params = params - lr * m_hat / (np.sqrt(v_hat) + eps)
    return new_params, {"m": m, "v": v, "t": t}

def adamw_step(params, grads, state, lr, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.0):
    # Adam-with-decoupled-weight-decay: same as adam_step but with weight_decay applied directly
    new_params, new_state = adam_step(params, grads, state, lr, beta1, beta2, eps)
    new_params = new_params - lr * weight_decay * params
    return new_params, new_state

Parte B — Compítelos en Rosenbrock¶

def run(opt_fn, lr, n_steps=2000, start=(-1.5, 2.0), **opt_kwargs):
    p = np.array(start, dtype=float)
    state = {}
    traj = [p.copy()]
    for _ in range(n_steps):
        g = rosenbrock_grad(p)
        p, state = opt_fn(p, g, state, lr, **opt_kwargs)
        traj.append(p.copy())
    return np.array(traj)

Ejecuta con tasas de aprendizaje ajustadas por optimizador (ya que tienen diferentes rangos de LR estable):

trajectories = {
    "SGD":      run(sgd_step,      lr=2e-3),
    "Momentum": run(momentum_step, lr=2e-3, beta=0.9),
    "Adam":     run(adam_step,     lr=5e-2),
    "AdamW":    run(adamw_step,    lr=5e-2, weight_decay=0.0),
}

(AdamW con weight_decay=0 es idéntico a Adam — esa es la comprobación.)

Parte C — Animar¶

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Contour
x_grid = np.linspace(-2, 2, 400)
y_grid = np.linspace(-1, 3, 400)
X, Y = np.meshgrid(x_grid, y_grid)
Z = (1 - X)**2 + 100 * (Y - X**2)**2

fig, ax = plt.subplots(figsize=(8, 6))
ax.contour(X, Y, Z, levels=np.logspace(-1, 3, 20), cmap="viridis", alpha=0.5)
ax.plot(1, 1, "r*", markersize=15, label="optimum")

lines = {name: ax.plot([], [], "-", label=name)[0] for name in trajectories}
ax.legend()
ax.set_xlim(-2, 2)
ax.set_ylim(-1, 3)

def init():
    for line in lines.values(): line.set_data([], [])
    return list(lines.values())

def update(frame):
    for name, line in lines.items():
        traj = trajectories[name]
        line.set_data(traj[:frame, 0], traj[:frame, 1])
    return list(lines.values())

anim = FuncAnimation(fig, update, init_func=init, frames=range(0, 2001, 20), interval=50, blit=True)
anim.save("rosenbrock-race.mp4", writer="ffmpeg")  # o .gif vía PillowWriter

Parte D — Interpretar¶

Escribe una interpretación de 200 palabras en experiments/04-optimizers-rosenbrock/INTERPRETATION.md. Aborda:

Camino de SGD. ¿Por qué oscila en la dirección empinada?
Camino de Momentum. Más suave, pero ¿hace overshoot? ¿Dónde?
Camino de Adam. ¿Por qué navega el valle estrecho tanto mejor? (Pista: normalización por parámetro vía sqrt(v̂).)
AdamW con weight_decay=0. Verifica que traza Adam exactamente. ¿Por qué es esta una comprobación útil?

Parte E — Una segunda barrida con `weight_decay > 0`¶

Ejecuta AdamW con weight_decay = 0.01 y weight_decay = 0.1. Añade a la gráfica. Esperado: el decay más fuerte tira de la trayectoria hacia el origen, ralentizando la convergencia a (1, 1). Anota.

Entregable¶

experiments/04-optimizers-rosenbrock/: - rosenbrock-race.mp4 (o .gif) — animación. - final_distances.json — para cada optimizador, la distancia desde la posición final hasta (1, 1). - INTERPRETATION.md — escrito de 200 palabras. - manifest.json — versiones, semillas.

Más un PNG estático de respaldo de las trayectorias finales (por si la animación no se renderiza en CI).

Aceptación¶

Trayectorias de los cuatro optimizadores renderizadas.
La distancia final de Adam a (1, 1) es < 0.1.
La distancia final de SGD es > que la de Adam (lo será — esa es la lección).
AdamW con weight_decay=0 coincide con Adam exactamente (diferencia elemento a elemento < 1e-12).
INTERPRETATION.md atribuye correctamente la ventaja de Adam al escalado por parámetro, no a "Adam es simplemente mejor".

Escollos¶

Usar la misma LR para los cuatro optimizadores. La LR segura de SGD es ~10× más pequeña que la de Adam en Rosenbrock. Si usas la LR de Adam para SGD, SGD diverge. Sintoniza por optimizador.
Olvidar t += 1 en Adam. La corrección de sesgo depende de t; si t = 0, divides por cero.
Usar AdamW = Adam + L2. Ese es el AdamW incorrecto. Desacoplado significa aplicado a los pesos directamente, no añadido al gradiente.
Fallos de renderizado de la animación. Puede que ffmpeg no esté instalado. Cae a PillowWriter para GIF.
Tratar "Adam gana" como lección universal. Adam gana aquí por el valle estrecho. En otros paisajes (bien condicionados, ruido bajo), SGD-con-momentum puede ganar. La lección es sobre el por qué, no sobre cuál.

Ampliación¶

Añade Nesterov momentum como quinta pista. Compara con momentum puro.
Grafica pérdida-vs-paso en un segundo panel. Adam debería tener una curva de descenso más suave.
Inicializa en un punto diferente (p. ej., (2, 2)) y vuelve a competir. Algunos patrones se mantienen; otros no.

Siguiente: 03-lr-schedules.md