Optimizers y Schedulers

🎯 Por qué importan los optimizadores

En el entrenamiento de redes neuronales, el objetivo es encontrar los parámetros \(\theta\) que minimizan una función de pérdida \(\mathcal{L}(\theta)\). El método más básico es el Stochastic Gradient Descent (SGD):

Donde \(\eta\) es la tasa de aprendizaje (learning rate) y \(\nabla_\theta \mathcal{L}\) es el gradiente de la pérdida respecto a los parámetros. Aunque conceptualmente simple, SGD vanilla tiene problemas serios:

• Oscilaciones en dimensiones con curvaturas muy diferentes (valles estrechos)
• Sensibilidad extrema a la elección del learning rate
• Se atasca en mínimos locales o puntos de silla
• Convergencia lenta cuando los gradientes son muy pequeños o ruidosos

La superficie de pérdida

Imagina la función de pérdida como un paisaje montañoso de alta dimensionalidad. SGD vanilla simplemente «baja la pendiente» en cada paso. Los problemas aparecen cuando:

• El valle es alargado y estrecho: SGD oscila de lado a lado en lugar de avanzar hacia el mínimo
• Hay puntos de silla (saddle points): el gradiente es ~0 pero no estamos en un mínimo
• La superficie tiene muchas mesetas donde el gradiente es casi nulo
• Diferentes parámetros necesitan learning rates distintos

Los optimizadores que veremos resuelven estos problemas de formas elegantes y complementarias.

La historia de los optimizadores en deep learning refleja un progreso fascinante. El SGD con mini-batches fue el caballo de batalla durante décadas, pero su sensibilidad al learning rate motivó la búsqueda de alternativas. En los años 2010, la comunidad desarrolló dos líneas paralelas: métodos basados en momentum (que acumulan información de gradientes pasados para suavizar la trayectoria) y métodos adaptativos (que ajustan el learning rate individualmente por parámetro). Adam (2015) fusionó ambas ideas y se convirtió en el optimizador por defecto. Más recientemente, AdamW (2019) corrigió un problema fundamental en la interacción entre Adam y weight decay, y se ha consolidado como el estándar para entrenar Transformers y modelos de gran escala.

🏎️ Momentum

La idea de momentum es inspirarse en la física: una pelota rodando cuesta abajo acumula velocidad. En lugar de depender solo del gradiente actual, momentum mantiene un «historial de velocidad» que suaviza las oscilaciones y acelera la convergencia.

Momentum clásico (Polyak, 1964)

Introducimos una variable de velocidad \(v_t\) que acumula gradientes exponencialmente:

El hiperparámetro \(\beta \in [0, 1)\) controla cuánto «recordamos» de los gradientes anteriores. Valores típicos: \(\beta = 0.9\) (estándar) o \(\beta = 0.99\) (alto momentum).

• Si \(\beta = 0\): es SGD vanilla (sin memoria)
• Si \(\beta = 0.9\): la velocidad es una media exponencial de los últimos ~10 gradientes
• Si \(\beta = 0.99\): promediamos los últimos ~100 gradientes

Intuitivamente, la velocidad \(v_t\) actúa como una media exponencial ponderada de los gradientes anteriores. Esto tiene un efecto doble: por un lado, suaviza el ruido estocástico inherente a los mini-batches (cada batch solo ve una fracción del dataset); por otro, permite acumular velocidad en direcciones consistentes, como una bola de nieve rodando cuesta abajo. Este principio, tomado directamente de la mecánica clásica, fue introducido por Boris Polyak en 1964 y sigue siendo uno de los trucos más efectivos en optimización numérica.

Momentum de Nesterov (NAG, 1983)

Yurii Nesterov propuso una mejora sutil pero poderosa: en lugar de calcular el gradiente en la posición actual \(\theta_t\), lo calculamos en la posición «hacia donde vamos» con el momentum actual. Es como «mirar hacia adelante» antes de dar el paso:

La diferencia clave está en \(\nabla_\theta \mathcal{L}(\theta_t \textcolor{#55efc4}{- \beta \cdot v_t})\): calculamos el gradiente en la posición anticipada, no en la posición actual. Esto proporciona una «corrección» que reduce el overshooting del momentum clásico.

Momentum clásico vs Nesterov

SGD con momentum (especialmente Nesterov) sigue siendo competitivo con optimizadores más modernos, sobre todo en visión por computador. Investigaciones como las de Wilson et al. (2017) mostraron que SGD con momentum, bien configurado con un scheduler adecuado, a menudo generaliza mejor que Adam en tareas de clasificación de imágenes, aunque Adam converge más rápido. Esta observación motivó la búsqueda de optimizadores que combinasen la velocidad de convergencia de Adam con la capacidad de generalización de SGD — un camino que eventualmente llevaría a AdamW.

⚡ Optimizadores adaptativos

La segunda gran familia de optimizadores son los adaptativos: ajustan automáticamente la tasa de aprendizaje para cada parámetro individual. La idea es simple: los parámetros que reciben gradientes grandes (frecuentes) deberían tener un learning rate más pequeño, y viceversa.

La motivación de los métodos adaptativos surge de una observación práctica: en una red neuronal típica, distintos parámetros pueden tener magnitudes de gradiente que difieren en varios órdenes de magnitud. Los pesos de las primeras capas, por ejemplo, suelen recibir gradientes mucho más pequeños que los de las últimas (el problema del vanishing gradient). Un learning rate único no puede satisfacer a ambos: si es grande para las primeras capas, es excesivo para las últimas, y viceversa. Los optimizadores adaptativos resuelven esto manteniendo estadísticas por parámetro de los gradientes pasados, y ajustando el step size en consecuencia.

AdaGrad (Duchi et al., 2011)

Adaptive Gradient fue el primer optimizador adaptativo. Acumula la suma de los cuadrados de todos los gradientes pasados:

Donde \(G_t\) es la suma acumulada de gradientes al cuadrado (por parámetro) y \(\epsilon \approx 10^{-8}\) evita la división por cero.

A pesar de este problema, AdaGrad introdujo una idea brillante que influenció a todos los optimizadores posteriores: adaptar el learning rate usando el historial de gradientes. La clave para resolver su problema de «learning rate decreciente» sería usar una ventana deslizante en lugar de una suma acumulativa — exactamente lo que haría RMSprop.

RMSprop (Hinton, 2012)

Geoffrey Hinton propuso RMSprop (Root Mean Square Propagation) en un slide de un curso, sin publicación formal. Resuelve el problema de AdaGrad usando una media exponencial en lugar de una suma acumulada:

Donde \(\rho\) (típicamente 0.9 o 0.99) controla la ventana de la media exponencial. Al «olvidar» gradientes antiguos, el learning rate no decae a cero.

RMSprop se puede entender como «AdaGrad con memoria limitada». Al usar una media exponencial en lugar de una suma acumulada, los gradientes recientes tienen más peso que los antiguos, lo que permite al optimizador adaptarse a cambios en la distribución de gradientes durante el entrenamiento. Curiosamente, RMSprop nunca se publicó en un paper formal — solo apareció en las diapositivas de la lección 6 del curso de redes neuronales de Hinton en Coursera (2012). A pesar de esto, se convirtió en uno de los optimizadores más utilizados, especialmente para RNNs.

AdaDelta (Zeiler, 2012)

AdaDelta es una extensión de RMSprop que elimina la necesidad de un learning rate inicial. Mantiene una media exponencial de los cuadrados de los deltas anteriores:

La clave es que el numerador usa la RMS de los updates anteriores, lo que actúa como un «learning rate implícito» que se adapta automáticamente.

Adam (Kingma & Ba, 2015)

Adaptive Moment Estimation combina lo mejor del momentum y de los métodos adaptativos. Mantiene dos medias exponenciales: el primer momento (media de gradientes = momentum) y el segundo momento (media de gradientes al cuadrado = adaptación):

Valores por defecto recomendados: \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon = 10^{-8}\).

AdamW (Loshchilov & Hutter, 2019)

AdamW es una corrección importante de Adam que implementa el weight decay de forma correcta. En Adam original, el weight decay se aplica al gradiente (L2 regularization), pero esto interactúa incorrectamente con la adaptación del learning rate. AdamW aplica el decay directamente a los pesos:

Donde \(\lambda\) es el coeficiente de weight decay (típicamente \(0.01\)). AdamW es el optimizador estándar de facto en visión por computador y NLP modernos (ViT, BERT, GPT, etc.).

La diferencia entre L2 regularization (Adam original) y decoupled weight decay (AdamW) puede parecer sutil, pero tiene consecuencias prácticas importantes. En Adam original, el término de regularización \(\lambda \theta\) se suma al gradiente antes de la normalización adaptativa, lo que significa que los parámetros con gradientes grandes reciben menos regularización efectiva — exactamente lo contrario de lo deseado. AdamW aplica el decay directamente a los pesos, independientemente de las estadísticas de gradientes. Loshchilov & Hutter (2019) demostraron que esta corrección mejora significativamente la generalización y hace que el weight decay interactúe correctamente con el learning rate schedule. En la práctica, esto significa que con AdamW puedes ajustar el learning rate y el weight decay de forma más independiente.

Otros optimizadores relevantes

La proliferación de optimizadores puede resultar abrumadora, pero en la práctica solo necesitas conocer bien unos pocos. La gran mayoría de proyectos modernos de deep learning usan AdamW (para Transformers y modelos generativos) o SGD con momentum (para CNNs en visión). Los demás optimizadores tienen nichos específicos: LAMB para pretraining distribuido con batch sizes masivos, Lion para reducir el consumo de memoria, y Adafactor para modelos de lenguaje enormes donde la memoria del optimizador es un cuello de botella. El siguiente widget te permite comparar visualmente cómo se comportan estos optimizadores en distintas superficies de pérdida.

📊 Comparación de optimizadores

La elección del optimizador está íntimamente ligada a la elección del scheduler de learning rate. Un optimizador como SGD con momentum necesita un scheduler agresivo (como Cosine Annealing o Step Decay) para alcanzar un buen rendimiento, mientras que AdamW funciona bien con warmup + cosine decay suave. La combinación optimizador + scheduler es a menudo más importante que la elección individual de cada uno. Además, ambas decisiones son excelentes candidatas para la optimización de hiperparámetros (HPO), donde herramientas como Optuna pueden explorar automáticamente las mejores combinaciones.

📈 Learning Rate Schedulers

Un scheduler (o policy de learning rate) modifica la tasa de aprendizaje \(\eta\) durante el entrenamiento. ¿Por qué? Porque el learning rate óptimo cambia a lo largo del entrenamiento:

• Al principio: un LR alto permite explorar rápidamente el espacio de parámetros
• Al final: un LR bajo permite «afinar» y converger a un mínimo preciso

La idea de variar el learning rate durante el entrenamiento tiene raíces profundas en la teoría de optimización estocástica. Las condiciones de Robbins-Monro (1951) establecen que para que SGD converja al óptimo global, el learning rate debe satisfacer \(\sum_{t=1}^{\infty} \eta_t = \infty\) (para explorar todo el espacio) y \(\sum_{t=1}^{\infty} \eta_t^2 < \infty\) (para que la varianza se reduzca a cero). Los schedulers modernos están diseñados para satisfacer estas condiciones de forma práctica, reduciendo el LR lo suficientemente rápido para converger, pero no tan rápido como para detener el aprendizaje prematuramente.

Step Decay

El más simple: reduce el LR por un factor fijo cada cierto número de epochs.

Donde \(\gamma\) es el factor de decaimiento (ej: 0.1) y \(S\) es el step size en epochs (ej: 30). Se suele configurar como: «reduce ×10 cada 30 epochs» (usado en ResNet original).

Exponential Decay

Reduce el LR exponencialmente en cada epoch:

Más suave que Step Decay. Típicamente \(\gamma \in [0.95, 0.99]\). Reduce de forma continua sin los «saltos» bruscos del Step Decay.

Cosine Annealing

Sigue una curva coseno para reducir el LR desde \(\eta_{\max}\) hasta \(\eta_{\min}\):

Donde \(T\) es el número total de epochs. Es el scheduler más popular en investigación moderna. Produce una transición suave que pasa la mayor parte del tiempo con LR intermedio.

Cosine Annealing con Warm Restarts (SGDR)

Extiende el cosine annealing con «restarts» periódicos que reinician el LR al valor máximo. Cada ciclo puede ser más largo que el anterior (multiplicando por \(T_{\text{mult}}\)):

Los restarts ayudan a escapar de mínimos locales y pueden usarse para hacer snapshot ensembles (guardar el modelo al final de cada ciclo).

Linear Warmup + Decay

Usado en Transformers (BERT, GPT): durante los primeros pasos, el LR crece linealmente desde 0 hasta \(\eta_{\max}\) (warmup), y después decae (lineal o coseno):

OneCycleLR (Smith, 2018)

La política «super-convergencia»: un único ciclo de coseno que sube el LR hasta un máximo y luego lo reduce a un valor muy bajo. También varía el momentum inversamente al LR:

• Fase 1 (warmup): LR sube de \(\eta_{\max}/\text{div}\) a \(\eta_{\max}\), momentum baja de 0.95 a 0.85
• Fase 2 (annealing): LR baja de \(\eta_{\max}\) a \(\eta_{\max}/(\text{div} \times \text{final\_div})\), momentum sube a 0.95

Permite entrenar modelos con batch sizes y LR mucho más altos, reduciendo significativamente el tiempo de entrenamiento.

ReduceLROnPlateau

A diferencia de los otros schedulers (basados en el tiempo), este es reactivo: reduce el LR cuando una métrica (ej: validation loss) deja de mejorar durante patience epochs.

• Ventaja: se adapta al comportamiento real del entrenamiento
• Desventaja: no se puede precomputar el schedule, dificulta la reproducibilidad
• Uso típico: ReduceLROnPlateau(patience=10, factor=0.1)

En la práctica, la elección del scheduler depende tanto del optimizador como del dominio de aplicación. La combinación más extendida en la investigación actual es warmup lineal + cosine decay, utilizada en la mayoría de papers de NLP (BERT, GPT, LLaMA) y visión (ViT, DINO, MAE). OneCycleLR es preferido cuando se busca entrenar rápido con batch sizes grandes, como demostró Leslie Smith (2018) en su trabajo sobre super-convergencia. ReduceLROnPlateau es útil en fases de exploración cuando no se conoce bien el comportamiento del modelo, pero su naturaleza reactiva lo hace menos reproducible que los schedulers predefinidos. El widget siguiente te permite visualizar y comparar las curvas de todos estos schedulers.

🔄 Comparación de schedulers

Estas combinaciones no son arbitrarias: reflejan años de experiencia empírica y teórica de la comunidad. En visión, SGD con momentum generaliza mejor que Adam en CNNs (Wilson et al., 2017), pero AdamW es superior para Vision Transformers (Dosovitskiy et al., 2021). En NLP, AdamW con warmup es casi universal desde BERT (Devlin et al., 2019). La lección principal es que no existe un «mejor optimizador universal»: la elección óptima depende de la arquitectura, el dataset y el presupuesto computacional.

🛠️ Ejemplo completo: Adam + Cosine Scheduler

Veamos un ejemplo completo y realista: entrenar una red para clasificación de imágenes usando AdamW como optimizador y Cosine Annealing con Warmup como scheduler. Esta es una configuración muy común en producción y en papers de investigación.

Configuración típica

Implementación

A continuación puedes ver la implementación completa en los dos frameworks más populares. Despliega cada bloque para ver el código:

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# ── Hiperparámetros ──────────────────────────────────────
EPOCHS = 100
WARMUP_EPOCHS = 5
LR = 3e-4
LR_MIN = 1e-6
WEIGHT_DECAY = 0.01
BATCH_SIZE = 128

# ── Dataset y DataLoader ─────────────────────────────────
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2470, 0.2435, 0.2616)),
])

train_ds = datasets.CIFAR10('./data', train=True,
                             download=True, transform=transform)
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE,
                          shuffle=True, num_workers=4)

val_ds = datasets.CIFAR10('./data', train=False,
                           transform=transforms.Compose([
                               transforms.ToTensor(),
                               transforms.Normalize((0.4914, 0.4822, 0.4465),
                                                    (0.2470, 0.2435, 0.2616)),
                           ]))
val_loader = DataLoader(val_ds, batch_size=256, shuffle=False)

# ── Modelo ────────────────────────────────────────────────
model = MyModel()  # Tu modelo (ResNet, ViT, etc.)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# ── Optimizador: AdamW ────────────────────────────────────
optimizer = AdamW(
    model.parameters(),
    lr=LR,
    betas=(0.9, 0.999),
    weight_decay=WEIGHT_DECAY,
)

# ── Scheduler: Linear Warmup + Cosine Annealing ──────────
warmup_scheduler = LinearLR(
    optimizer,
    start_factor=0.01,       # empieza en lr * 0.01
    end_factor=1.0,          # termina en lr * 1.0
    total_iters=WARMUP_EPOCHS,
)
cosine_scheduler = CosineAnnealingLR(
    optimizer,
    T_max=EPOCHS - WARMUP_EPOCHS,
    eta_min=LR_MIN,
)
scheduler = SequentialLR(
    optimizer,
    schedulers=[warmup_scheduler, cosine_scheduler],
    milestones=[WARMUP_EPOCHS],
)

# ── Función de pérdida ────────────────────────────────────
criterion = nn.CrossEntropyLoss()

# ── Bucle de entrenamiento ────────────────────────────────
for epoch in range(EPOCHS):
    model.train()
    total_loss = 0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    # Step del scheduler (una vez por epoch)
    scheduler.step()

    # Logging
    current_lr = optimizer.param_groups[0]['lr']
    avg_loss = total_loss / len(train_loader)
    print(f'Epoch {epoch+1}/{EPOCHS} | '
          f'Loss: {avg_loss:.4f} | '
          f'LR: {current_lr:.6f}')

    # ── Validación ─────────────────────────────────────────
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    acc = 100. * correct / total
    print(f'  Val Accuracy: {acc:.2f}%')

import tensorflow as tf
from tensorflow import keras
import numpy as np
import math

# ── Hiperparámetros ──────────────────────────────────────
EPOCHS = 100
WARMUP_EPOCHS = 5
LR = 3e-4
LR_MIN = 1e-6
WEIGHT_DECAY = 0.01
BATCH_SIZE = 128

# ── Dataset ───────────────────────────────────────────────
(x_train, y_train), (x_val, y_val) = keras.datasets.cifar10.load_data()
x_train = x_train.astype('float32') / 255.0
x_val = x_val.astype('float32') / 255.0

# Normalización con media y std de CIFAR-10
mean = np.array([0.4914, 0.4822, 0.4465])
std = np.array([0.2470, 0.2435, 0.2616])
x_train = (x_train - mean) / std
x_val = (x_val - mean) / std

STEPS_PER_EPOCH = len(x_train) // BATCH_SIZE
TOTAL_STEPS = STEPS_PER_EPOCH * EPOCHS
WARMUP_STEPS = STEPS_PER_EPOCH * WARMUP_EPOCHS

# ── Custom LR Schedule: Warmup + Cosine ──────────────────
class WarmupCosineSchedule(keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, lr_max, lr_min, warmup_steps, total_steps):
        super().__init__()
        self.lr_max = lr_max
        self.lr_min = lr_min
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps

    def __call__(self, step):
        step = tf.cast(step, tf.float32)
        # Warmup phase
        warmup_lr = self.lr_max * (step / self.warmup_steps)
        # Cosine decay phase
        progress = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
        progress = tf.clip_by_value(progress, 0.0, 1.0)
        cosine_lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
            1 + tf.cos(progress * math.pi)
        )
        return tf.where(step < self.warmup_steps, warmup_lr, cosine_lr)

    def get_config(self):
        return {"lr_max": self.lr_max, "lr_min": self.lr_min,
                "warmup_steps": self.warmup_steps,
                "total_steps": self.total_steps}

lr_schedule = WarmupCosineSchedule(
    lr_max=LR,
    lr_min=LR_MIN,
    warmup_steps=WARMUP_STEPS,
    total_steps=TOTAL_STEPS,
)

# ── Optimizador: AdamW con schedule ──────────────────────
optimizer = keras.optimizers.AdamW(
    learning_rate=lr_schedule,
    beta_1=0.9,
    beta_2=0.999,
    weight_decay=WEIGHT_DECAY,
)

# ── Modelo ────────────────────────────────────────────────
model = MyKerasModel()  # Tu modelo Keras

model.compile(
    optimizer=optimizer,
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'],
)

# ── Data augmentation ─────────────────────────────────────
data_augmentation = keras.Sequential([
    keras.layers.RandomFlip('horizontal'),
    keras.layers.RandomTranslation(0.1, 0.1),
])

# ── Entrenamiento ─────────────────────────────────────────
history = model.fit(
    data_augmentation(x_train),
    y_train,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    validation_data=(x_val, y_val),
    verbose=1,
)

# ── Evaluar ───────────────────────────────────────────────
val_loss, val_acc = model.evaluate(x_val, y_val, verbose=0)
print(f'Val Accuracy: {val_acc*100:.2f}%')

Si quieres ir más allá del ajuste manual y explorar automáticamente las mejores combinaciones de optimizador, scheduler y sus hiperparámetros, consulta nuestro submódulo de Optimización de Hiperparámetros (HPO), donde cubrimos desde Random Search hasta optimización bayesiana con Optuna.