El promedio de pesos estocástico (SWA, por sus siglas en inglés) es una técnica de promediado de modelos en el Deep learning que mejora la generalización al promediar los pesos de una Neural network en múltiples puntos a lo largo de su trayectoria de entrenamiento. A diferencia del entrenamiento tradicional que conserva solo los pesos finales, SWA recoge pesos a intervalos regulares durante la fase posterior del entrenamiento y calcula su media aritmética. Este procedimiento simple a menudo conduce a soluciones más planas en el paisaje de pérdida, que se asocian con un mejor rendimiento de prueba y robustez. SWA fue introducido por Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov y Andrew Gordon Wilson en su artículo de 2018 "Averaging Weights Leads to Wider Optima and Better Generalization". El método no requiere cambios en la arquitectura del modelo subyacente ni en la función de pérdida, lo que facilita su integración en los pipelines de entrenamiento existentes.
La idea central detrás de SWA es que durante el entrenamiento con descenso de gradiente estocástico (SGD) o sus variantes, los parámetros del modelo oscilan alrededor de una región de baja pérdida. Promediar estos parámetros puede suavizar las oscilaciones y producir una solución que se encuentra en un mínimo amplio y plano. Se hipotetiza que estos mínimos planos generalizan mejor que los mínimos agudos porque son menos sensibles a las perturbaciones y a los cambios entre las distribuciones de entrenamiento y prueba. Se ha demostrado que SWA mejora la precisión y la calibración en diversas arquitecturas y tareas, incluyendo clasificación de imágenes, detección de objetos y modelado de lenguaje.
Mecanismo y Algoritmo
SWA opera en dos fases. En la primera fase, el modelo se entrena con un programa de tasa de aprendizaje estándar (por ejemplo, annealing coseno o decaimiento por pasos) para alcanzar una región cercana a un mínimo. En la segunda fase, el entrenamiento continúa con una tasa de aprendizaje constante o cíclica, y los pesos se recogen al final de cada época o después de un número fijo de iteraciones. Los pesos de SWA se actualizan como un promedio móvil: n_swa = n_swa + 1; w_swa = (w_swa * (n_swa - 1) + w) / n_swa, donde w son los pesos actuales y w_swa son los pesos promediados. El modelo final utiliza w_swa.
Un detalle clave es que las estadísticas de normalización por lotes deben recalcularse después del promediado, porque la media móvil y la varianza de las capas de normalización por lotes no se promedian junto con los pesos. En la práctica, después de calcular los pesos de SWA, se realiza un pase hacia adelante sobre los datos de entrenamiento para actualizar las estadísticas de normalización por lotes. Este paso es crucial para los modelos que utilizan normalización por lotes, como ResNets.
Motivación Teórica
El éxito de SWA a menudo se explica a través de la lente de la geometría del paisaje de pérdida. La investigación de Izmailov et al. y trabajos posteriores han demostrado que SGD tiende a converger a puntos cercanos al límite de una región de baja pérdida, mientras que promediar múltiples puntos mueve la solución hacia el centro de esa región. Este punto central típicamente se encuentra en una cuenca más plana, que se asocia con una mejor generalización. La conexión entre mínimos planos y generalización ha sido estudiada extensamente, con vínculos con los límites de PAC-Bayes y la agudeza de la función de pérdida.
Otra perspectiva relaciona SWA con la inferencia bayesiana. Bajo ciertas suposiciones, la trayectoria de SGD puede verse como una cadena de Markov que muestrea de una distribución posterior. Promediar a lo largo de esta trayectoria aproxima la media posterior, que es la predicción óptima de Bayes para el error de pérdida cuadrática. Esta interpretación vincula SWA con los métodos de Monte Carlo de cadena de Markov de gradiente estocástico, aunque SWA es más simple y no requiere un ajuste cuidadoso de las escalas de ruido.
Variantes y Extensiones
Se han propuesto varias variantes de SWA para mejorar aún más el rendimiento. Una extensión notable es el Promedio de Pesos Estocástico Gaussiano (SWAG, por sus siglas en inglés), introducido por Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov y Andrew Gordon Wilson en 2019. SWAG aproxima la distribución posterior de los pesos ajustando una gaussiana utilizando los primeros y segundos momentos de los pesos recogidos. Esto permite la estimación de incertidumbre y una mejor calibración, y puede utilizarse para el aprendizaje profundo bayesiano.
Otra variante es el Ensamblado Geométrico Rápido (FGE, por sus siglas en inglés), que utiliza un programa de tasa de aprendizaje cíclico para recoger pesos de diferentes modos del paisaje de pérdida. SWA puede combinarse con FGE para lograr un rendimiento aún mejor. Además, SWA se ha integrado con técnicas como programas de tasa de aprendizaje, aumento de datos y poda para mejorar la robustez y la eficiencia.
Aplicaciones
SWA se ha aplicado con éxito en una amplia gama de dominios. En visión por computadora, mejora la precisión de los clasificadores de imágenes en conjuntos de datos como CIFAR-10, CIFAR-100 e ImageNet. Por ejemplo, se ha demostrado que SWA aumenta la precisión de prueba de un PreAct ResNet-164 en CIFAR-100 del 81.15% al 82.90% sin cambios arquitectónicos. En detección de objetos, SWA ayuda a estabilizar el entrenamiento y mejorar la precisión media promedio.
En procesamiento de lenguaje natural, SWA se ha utilizado para entrenar Transformers y modelos de lenguaje grandes de manera más efectiva. Se ha demostrado que reduce el sobreajuste y mejora la perplejidad en tareas de modelado de lenguaje. SWA también beneficia a modelos generativos como GANs y modelos de difusión al proporcionar actualizaciones de pesos más suaves, lo que puede conducir a un entrenamiento más estable y muestras de mayor calidad.
Más allá del aprendizaje supervisado estándar, SWA se ha adaptado para el aprendizaje semi-supervisado, el aprendizaje por transferencia y la adaptación de dominio. Su simplicidad y generalidad lo convierten en una herramienta valiosa en el kit de herramientas del practicante de aprendizaje automático.
Comparación con Otras Técnicas
SWA a menudo se compara con otros métodos de regularización y ensamblado. A diferencia de los ensamblados tradicionales que promedian las predicciones de múltiples modelos entrenados de forma independiente, SWA promedia los pesos, lo que es computacionalmente más barato y requiere solo una ejecución de entrenamiento. Esto lo hace particularmente atractivo cuando el almacenamiento del modelo o el costo de inferencia es una preocupación.
En comparación con técnicas como Dropout o normalización por lotes, SWA no modifica la dinámica de entrenamiento sino que procesa posteriormente los pesos. Puede utilizarse junto con estos métodos. SWA también difiere de Adam y otros optimizadores adaptativos, que mantienen tasas de aprendizaje por parámetro; SWA típicamente se aplica sobre trayectorias de SGD o Adam.
Una limitación de SWA es que requiere un programa para determinar cuándo comenzar a recoger pesos. Si la recolección comienza demasiado temprano, el promedio puede estar dominado por pesos subóptimos; si es demasiado tarde, los beneficios disminuyen. Sin embargo, en la práctica, una regla simple es comenzar a recoger después de que la tasa de aprendizaje haya decaído a un valor bajo.
Consideraciones Prácticas
Implementar SWA es sencillo. La mayoría de los marcos de aprendizaje profundo, como PyTorch y TensorFlow, proporcionan utilidades o ejemplos para SWA. Los pasos principales son: (1) entrenar el modelo con un programa estándar, (2) continuar el entrenamiento con una tasa de aprendizaje constante o cíclica, (3) recoger pesos a intervalos, (4) calcular el promedio móvil y (5) recalcular las estadísticas de normalización por lotes.
Se ha demostrado que SWA es robusto a las elecciones de hiperparámetros. La tasa de aprendizaje para la segunda fase típicamente se establece en el valor mínimo del programa inicial o en una constante pequeña. La frecuencia de recolección de pesos puede ser cada época o cada pocas cientos de iteraciones. En la práctica, recoger pesos durante 5 a 20 épocas suele ser suficiente.
Impacto y Recepción
La introducción de SWA ha tenido un impacto significativo en la comunidad de aprendizaje profundo. Proporcionó un método simple pero efectivo para mejorar la generalización sin costo computacional adicional durante la inferencia. El artículo ha sido ampliamente citado y ha inspirado numerosos trabajos de seguimiento sobre promediado de pesos y mínimos planos. SWA es ahora una técnica estándar en muchos pipelines de entrenamiento, especialmente en investigación académica y competiciones.
SWA también ha influido en el desarrollo de otros métodos de promediado, como el Promedio Móvil Exponencial (EMA, por sus siglas en inglés), que se utiliza comúnmente en el entrenamiento de modelos generativos y modelos de lenguaje grandes. EMA es una variante que da más peso a los pesos recientes, y a menudo se utiliza en la práctica debido a su simplicidad y efectividad. SWA y EMA son complementarios, y algunos marcos permiten cambiar entre ellos.
Limitaciones y Direcciones Futuras
A pesar de sus beneficios, SWA tiene limitaciones. Asume que el paisaje de pérdida es relativamente suave y que promediar pesos es significativo, lo que puede no cumplirse para todas las arquitecturas o tareas. Por ejemplo, los modelos con normalización por lotes requieren un manejo cuidadoso, y los modelos con espacios de parámetros complejos (por ejemplo, redes recurrentes) pueden no beneficiarse tanto. Además, SWA no proporciona estimaciones de incertidumbre a menos que se extienda (por ejemplo, SWAG).
Las direcciones de investigación futura incluyen el desarrollo de programas de promediado adaptativos, la combinación de SWA con métodos bayesianos y la aplicación de SWA a nuevos dominios como el aprendizaje por refuerzo y el aprendizaje federado. A medida que el aprendizaje profundo continúa evolucionando, las técnicas de promediado de pesos como SWA siguen siendo una herramienta fundamental para mejorar el rendimiento y la fiabilidad de los modelos.