Stochastic Gradient Descent (SGD)

El descenso estocástico del gradiente (SGD) es un algoritmo de optimización que entrena un modelo actualizando repetidamente sus pesos mediante el gradiente de la pérdida calculado sobre una pequeña muestra aleatoria de los datos de entrenamiento, en lugar de sobre el conjunto de datos completo. El gradiente de cada muestra es una estimación ruidosa pero insesgada del gradiente del conjunto completo, por lo que SGD realiza muchas actualizaciones baratas en el tiempo que el descenso del gradiente por lotes completos tardaría en hacer una sola actualización costosa. El método deriva del método de aproximación estocástica de Robbins y Monro de 1951, y con momento sigue siendo un optimizador estándar para las redes neuronales convolucionales (CNN).

Cómo funciona SGD#

Cada paso de SGD toma un minib lote B, calcula el gradiente de la pérdida correspondiente mediante retropropagación y mueve cada peso un pequeño paso en sentido contrario a ese gradiente: w ← w − η·∇L_B(w), donde η es la tasa de aprendizaje. En sentido estricto, SGD utiliza un ejemplo por paso, pero los frameworks de aprendizaje profundo lo aplican a minib lotes de decenas a miles de ejemplos; en la práctica, «SGD» y «descenso del gradiente por minib lotes» son sinónimos. Un tamaño de lote mayor produce un gradiente menos ruidoso y aprovecha mejor el hardware, a costa de realizar menos actualizaciones por época.

El entrenamiento mezcla aleatoriamente los datos al inicio de cada época, los recorre en minib lotes y da un paso por lote. El ajuste de una recta a 1.000 puntos ruidosos muestra todo el proceso:

import torch

torch.manual_seed(0)

# 1,000 noisy samples of y = 3x + 2
x = torch.randn(1000, 1)
y = 3 * x + 2 + 0.1 * torch.randn(1000, 1)

model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9)

for epoch in range(5):
    # Reshuffle every epoch, then update the weights once per 32-sample mini-batch
    for idx in torch.randperm(1000).split(32):
        loss = torch.nn.functional.mse_loss(model(x[idx]), y[idx])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f"epoch {epoch}: w = {model.weight.item():.3f}, b = {model.bias.item():.3f}, loss = {loss.item():.4f}")

Tras dos épocas, el peso llega a 3.081 y el sesgo a 2.012, con 64 actualizaciones en total. Durante las tres épocas siguientes, el peso sigue variando entre 2.967 y 3.003 en vez de estabilizarse, porque cada gradiente de minib lote contiene ruido de muestreo, y la pérdida del último lote se mantiene cerca de 0.01, la varianza del ruido añadido a y.

Tasa de aprendizaje, ruido y tamaño del lote#

Con una tasa de aprendizaje constante, SGD fluctúa en una región alrededor del mínimo cuyo tamaño aumenta con la tasa de aprendizaje y el ruido del gradiente. La teoría de la aproximación estocástica solo garantiza la convergencia cuando la tasa de aprendizaje disminuye hasta cero, por lo que en la práctica se reduce con el tiempo mediante programaciones escalonadas o coseno, normalmente después de un breve calentamiento. El ruido del gradiente de los lotes pequeños también guía SGD hacia mínimos planos que generalizan mejor que los mínimos pronunciados que suele encontrar el entrenamiento con lotes grandes, una forma implícita de regularización.

El tamaño del lote y la tasa de aprendizaje interactúan. Aumentar la tasa de aprendizaje en proporción al tamaño del lote, con un calentamiento de cinco épocas, permitió a ResNet-50 entrenarse con ImageNet en una hora usando minilotes de 8192 imágenes y sin perder precisión.

Momentum y momentum de Nesterov#

En un valle largo y estrecho de la superficie de pérdida, el gradiente apunta sobre todo de un lado a otro del valle, por lo que SGD básico zigzaguea entre las paredes mientras avanza lentamente por el fondo. Momentum mantiene una velocidad acumulada, v ← μ·v + g, y avanza siguiendo esa velocidad, w ← w − η·v; así, las direcciones coherentes del gradiente acumulan velocidad y las direcciones oscilantes se cancelan. El valor predeterminado habitual del coeficiente es μ = 0,9. El momentum de Nesterov evalúa el gradiente después de un paso de anticipación siguiendo la velocidad, lo que amortigua aún más el sobrepasamiento. SGD de PyTorch implementa ambos y aplica weight_decay añadiendo λ·w al gradiente, lo que en SGD equivale a una penalización L2 sobre los pesos.

SGD frente a otros optimizadores#

• Descenso del gradiente por lotes: Calcula el gradiente exacto sobre todo el conjunto de datos en cada actualización, por lo que cada época produce un único paso y las necesidades de memoria aumentan con el conjunto de datos.
• Adam: Asigna a cada parámetro su propio tamaño de paso a partir de las medias móviles de sus gradientes y de los gradientes al cuadrado, y almacena dos valores de estado por parámetro, frente a uno en SGD con momentum. Adam avanza más rápido al principio y es menos sensible a la tasa de aprendizaje, mientras que SGD con momentum bien ajustado suele igualar o superar su precisión final en las CNN. AdamW, su variante con decaimiento de pesos desacoplado, es el optimizador estándar para los transformers.
• Muon: Ortogonaliza en conjunto la actualización de momentum de cada matriz de pesos y se ha escalado para entrenar modelos de lenguaje de gran tamaño.

El entrenamiento de Ultralytics YOLO crea un SGD de PyTorch con momento de Nesterov cuando se establece optimizer="SGD", utilizando lr0=0.01 y momentum=0.937 de forma predeterminada. El valor predeterminado de optimizer="auto" ignora lr0 y elige MuSGD, un método híbrido que combina actualizaciones de SGD y Muon, para ejecuciones de más de 10.000 iteraciones, y AdamW para las más cortas.

Dónde se usa SGD#

SGD con momentum entrenó la mayoría de las CNN emblemáticas de ImageNet. Los modelos ResNet originales usaban momentum de 0,9, decaimiento de pesos de 0,0001 y minilotes de 256 imágenes, con una tasa de aprendizaje que empezaba en 0,1 y se dividía por 10 cada vez que el error se estancaba. Fuera del aprendizaje profundo, SGDClassifier y SGDRegressor de scikit-learn ajustan modelos lineales como la regresión logística y las máquinas de vectores de soporte lineales con SGD, y su método partial_fit entrena con conjuntos de datos demasiado grandes para caber en memoria, procesándolos fragmento a fragmento.