English · Español

Lab 02 — `SGD` (con momentum) y `Adam`¶

Objetivo: implementa dos optimizadores desde cero — SGD con momentum opcional, y Adam con corrección de sesgo — y contrasta la trayectoria de Adam contra torch.optim.Adam en una cuadrática de juguete. ~70 LOC por Borja.

Tiempo estimado: 120–150 minutos.

Prereqs: Lab 00 + Lab 01 cerrados. Teoría 03 leída.

🇪🇸 La matemática se derivó en Phase 4; aquí solo la encarnamos en step() y zero_grad(). Lo único técnico es la corrección de sesgo de Adam, que es exactamente la suma de la serie geométrica 1 + β + β² + .... El cross-check contra torch.optim.Adam te dice si tu fórmula está en off-by-one — un error muy fácil de cometer en el contador t.

Qué produces¶

src/minimodel/optim.py — clase base Optimizer, SGD, Adam.
tests/test_optimizers.py — tests de convergencia y el cross-check con PyTorch.

Referencia matemática (pégala primero en tu journal)¶

SGD (sin momentum): θ ← θ - η · g.

SGD con momentum (la convención de PyTorch, que es la que replicamos):

v ← β · v + (1 - β) · g            # media móvil exponencial de g
θ ← θ - η · v

La forma "amortiguada" de arriba es la que usa PyTorch cuando dampening != 0. Con dampening=0 (por defecto), la actualización es v ← β · v + g; θ ← θ - η · v. Elige la forma amortiguada ((1-β)·g) para que la interpretación como media móvil quede limpia; documenta la divergencia con el default de PyTorch.

Adam:

t ← t + 1                                       # 1-indexed; la PRIMERA actualización fija t=1
m ← β₁ · m + (1 - β₁) · g
v ← β₂ · v + (1 - β₂) · g²
m̂ ← m / (1 - β₁^t)                              # corrección de sesgo
v̂ ← v / (1 - β₂^t)
θ ← θ - η · m̂ / (√v̂ + ε)

Defaults: β₁ = 0.9, β₂ = 0.999, ε = 1e-8, lr = 1e-3.

AdamW NO es este lab. La distinción weight-decay-vs-L2 se menciona aquí para que la conozcas; la Phase 10 la implementa después de derivar por qué L2 dentro del gradiente interactúa mal con el precondicionamiento de Adam.

TODOs¶

Bloque A — clase base `Optimizer`¶

En src/minimodel/optim.py:

from typing import Iterable
import numpy as np
from minimodel.nn.module import Parameter


class Optimizer:
    """Base class. Materializes params as a list; per-parameter state keyed by id(p)."""

    def __init__(self, params: Iterable[Parameter], lr: float) -> None:
        # TODO: self.params = list(params). NOTE: iterables exhaust.
        # TODO: self.lr = lr.
        # TODO: self.state: dict[int, dict[str, np.ndarray]] = {id(p): {} for p in self.params}
        raise NotImplementedError

    def step(self) -> None:
        raise NotImplementedError

    def zero_grad(self) -> None:
        # TODO: for p in self.params: p.grad = None.
        raise NotImplementedError

Bloque B — `SGD`¶

class SGD(Optimizer):
    def __init__(
        self,
        params: Iterable[Parameter],
        lr: float,
        momentum: float = 0.0,
    ) -> None:
        super().__init__(params, lr)
        self.momentum = momentum
        if momentum > 0:
            # TODO: for each param, initialize state["velocity"] = np.zeros_like(p.data).
            raise NotImplementedError

    def step(self) -> None:
        # TODO: for each param p:
        #   if p.grad is None: continue
        #   g = p.grad
        #   if self.momentum > 0:
        #       v = self.state[id(p)]["velocity"]
        #       v = self.momentum * v + (1 - self.momentum) * g     # dampened form
        #       self.state[id(p)]["velocity"] = v
        #       update = v
        #   else:
        #       update = g
        #   p.data = p.data - self.lr * update
        raise NotImplementedError

p.data = p.data - self.lr * update es una op out-of-place. NO uses p.data -= ...: la mutación in-place puede confundir a autograd si alguna vista de tensor sigue referenciando el buffer antiguo.

Bloque C — `Adam`¶

class Adam(Optimizer):
    def __init__(
        self,
        params: Iterable[Parameter],
        lr: float = 1e-3,
        betas: tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
    ) -> None:
        super().__init__(params, lr)
        self.beta1, self.beta2 = betas
        self.eps = eps
        # TODO: for each param, initialize state["m"] = zeros_like, state["v"] = zeros_like,
        #       state["t"] = 0 (the per-parameter step counter).
        raise NotImplementedError

    def step(self) -> None:
        # TODO: for each param p:
        #   if p.grad is None: continue
        #   st = self.state[id(p)]
        #   st["t"] += 1                                   # 1-indexed; FIRST step is t=1
        #   t = st["t"]
        #   g = p.grad
        #   st["m"] = self.beta1 * st["m"] + (1 - self.beta1) * g
        #   st["v"] = self.beta2 * st["v"] + (1 - self.beta2) * (g * g)
        #   m_hat = st["m"] / (1 - self.beta1 ** t)
        #   v_hat = st["v"] / (1 - self.beta2 ** t)
        #   p.data = p.data - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
        raise NotImplementedError

Trampa off-by-one. La primerísima llamada a step() debe usar t=1, no t=0. Con t=0, 1 - β^0 = 0 y divides por cero.

Tests¶

En tests/test_optimizers.py:

Bloque D — convergencia de `SGD`¶

test_sgd_reduces_quadratic: Optimiza f(θ) = (θ - 3)² desde θ₀ = 0 con lr=0.1 durante 100 pasos. Asegura que |θ - 3| < 1e-3 al final y que la pérdida es monótonamente no creciente.
test_sgd_momentum_faster_than_plain: Misma cuadrática. Compara SGD plano (momentum=0) con SGD momentum=0.9. Tras 30 pasos, el SGD con momentum debe estar más cerca del óptimo (desigualdad estricta). Es un test de regresión sobre la implementación, no un teorema profundo.
test_sgd_zero_grad_resets: Fija manualmente el grad de un parámetro, llama a zero_grad, asegura que grad is None.

Bloque E — convergencia de `Adam`¶

test_adam_reduces_quadratic: La misma f(θ) = (θ - 3)² desde θ₀ = 0. Con lr=0.1 y las betas por defecto, Adam alcanza |θ - 3| < 1e-2 en 200 pasos.
test_adam_first_step_uses_t_eq_1: Sanity-check de la corrección de sesgo: con lr=0.1, tras UN paso sobre f(θ) = θ² con θ₀=1, la magnitud de la actualización es ≈ lr (porque m̂ y √v̂ son ambos ≈ |g| tras la corrección de sesgo). Asegura que |θ - (1 - 0.1)| < 1e-6. Con t=0 (el bug off-by-one), este test produce nan o valores ridículamente erróneos.

Bloque F — cross-check con PyTorch¶

Restricción: PyTorch entra únicamente en este fichero de test. NO entra en src/minimodel/. El cross-check existe para verificar nuestras fórmulas, no para importar el framework a la librería.

test_adam_matches_torch_on_quadratic:

import torch  # test fixture only

np.random.seed(0)
torch.manual_seed(0)

# Toy quadratic: minimize ||θ - target||² over a 5-dim θ.
target_np = np.random.randn(5)

# OUR Adam.
theta = Parameter(np.zeros(5))
our_opt = Adam([theta], lr=1e-2)
for _ in range(100):
    our_opt.zero_grad()
    # loss = sum((theta - target)²). Use minitorch ops so .backward() populates theta.grad.
    # TODO: build the loss tensor, call loss.backward(), then our_opt.step().
    ...

# PyTorch Adam.
theta_t = torch.zeros(5, requires_grad=True)
target_t = torch.tensor(target_np)
torch_opt = torch.optim.Adam([theta_t], lr=1e-2)
for _ in range(100):
    torch_opt.zero_grad()
    loss_t = ((theta_t - target_t) ** 2).sum()
    loss_t.backward()
    torch_opt.step()

assert np.allclose(theta.data, theta_t.detach().numpy(), atol=1e-5)

Tolerancia 1e-5 sobre 100 pasos. Si el test falla por 1e-2, sospecha del off-by-one en t. Si falla por 1e-6 solo tras muchos pasos, sospecha de un eps ausente o una β por defecto incorrecta.

Bloque G — casos extremos¶

test_optimizer_with_none_grad_skips: Un parámetro con grad = None no debe lanzar; step() simplemente lo salta.
test_optimizer_state_isolated_per_parameter: Dos parámetros de la misma shape NO deben compartir state["m"] ni state["velocity"]. Verifícalo con id(...) y mutando uno y comprobando el otro.
test_optimizer_does_not_track_late_added_params: Construye un Adam([p1]), después crea p2 y asegura que opt.state no tiene clave para id(p2). Documenta: los parámetros añadidos tras construir el optimizador no se rastrean. PyTorch tiene el mismo comportamiento.

Restricciones¶

Nada de PyTorch en src/minimodel/. PyTorch solo se permite en tests/test_optimizers.py, como oráculo de referencia.
Nada de AdamW, nada de schedulers de learning rate, nada de gradient clipping. Phase 10 + Phase 18.
SGD no implementa Nesterov momentum. Menciónalo en un comentario; no lo implementes.
Ámbito A13: los optimizadores son independientes del dominio; nada de código específico de verbos en este lab.

Escollos¶

Off-by-one en t. El primer step() debe usar t=1. Testéalo (Bloque E).
AdamW vs Adam + L2 weight decay. Adam con decay L2 mete λθ dentro del gradiente antes del precondicionamiento, lo que hace que el decay efectivo no sea constante. AdamW aplica θ ← θ - lr·λ·θ fuera del precondicionamiento. No implementamos ninguna de las dos formas de weight-decay en este lab — solo debes saber que son distintas. La Phase 10 cierra el hueco.
Aliasing de Parameter compartido. Dos atributos del modelo apuntando al mismo Parameter (Lab 00, Bloque F) acaban en self.params dos veces. El optimizador actualiza el parámetro dos veces por step. Está mal para embeddings atados (la Phase 17 deduplicará). Por ahora: documenta, no arregles.
state indexado por id(p) vs por índice. Las claves de id sobreviven cuando self.params se reordena (no se reordena, pero el invariante importa); las claves de índice se rompen. Quédate con id.
np.sqrt de un v̂ cero o casi cero. El + eps está dentro de la suma del sqrt: lr · m̂ / (√v̂ + ε). La referencia de PyTorch lo calcula así; igualar el parentizado hasta 1e-5 requiere igualar la posición de eps.
p.grad lo pone a None zero_grad, no a ceros. Convención de PyTorch. Testéalo.

Condiciones de parada¶

Terminado cuando:

SGD ocupa ≤ 30 líneas, Adam ocupa ≤ 40 líneas.
Todos los tests de los Bloques D–G en verde, incluido el cross-check con PyTorch a 1e-5.
mypy --strict src/minimodel/optim.py limpio.
ruff check src/minimodel/optim.py limpio.
Puedes derivar el factor de corrección de sesgo 1 - β^t a partir de la serie geométrica (una línea en tu journal).
Puedes explicar en una frase por qué AdamW no es lo mismo que Adam + L2.

Cuándo consultar `solutions/`¶

Tras pasar todos los tests. solutions/02-optimizers-ref.md (en la apertura de la fase) recorre las formulaciones alternativas (SGD no amortiguado de PyTorch, Nesterov momentum, AdamW) y dónde encajarían.

Siguiente lab: lab/03-train-tense-mlp.md.

Lab 02 — SGD (con momentum) y Adam¶