AdamW es un algoritmo de optimización utilizado en el Machine learning y el Deep learning para entrenar redes neuronales. Es una variante del optimizador Adam que desacopla la disminución de peso de las actualizaciones adaptativas de la tasa de aprendizaje. El algoritmo fue introducido en 2017 por Ilya Loshchilov y Frank Hutter, y se ha convertido en una opción estándar para entrenar modelos grandes, incluidos los transformers y los grandes modelos de lenguaje. Al separar la disminución de peso de las actualizaciones basadas en gradientes, AdamW aborda un problema conocido en el optimizador Adam original, donde la disminución de peso se aplicaba de una manera que interfería con las tasas de aprendizaje adaptativas, lo que conducía a una convergencia y generalización subóptimas.
La motivación principal para AdamW surgió de la observación de que en Adam, el término de regularización L2 (a menudo utilizado como disminución de peso) se divide por la raíz cuadrada del promedio móvil exponencial de los gradientes al cuadrado. Este acoplamiento hace que la disminución de peso efectiva varíe entre parámetros y a lo largo del tiempo, lo que puede dificultar la optimización. AdamW propone una solución sencilla: aplicar la disminución de peso directamente a los parámetros después de la actualización del gradiente, independientemente del escalado adaptativo de la tasa de aprendizaje. Se ha demostrado que este desacoplamiento mejora el rendimiento del entrenamiento en una variedad de tareas, incluida la clasificación de imágenes y el modelado de lenguaje.
Antecedentes: Optimización en el aprendizaje automático
Entrenar una red neuronal implica minimizar una función de pérdida, típicamente utilizando variantes del descenso de gradiente estocástico (SGD). En SGD, los parámetros del modelo se actualizan iterativamente moviéndose en la dirección del gradiente negativo de la pérdida, calculado sobre un subconjunto aleatorio de los datos de entrenamiento. El tamaño del paso, o tasa de aprendizaje, controla cuán grande es cada actualización. A lo largo de los años, se han propuesto muchas mejoras para acelerar la convergencia y mejorar el rendimiento final, como el momento, las tasas de aprendizaje adaptativas y la disminución de peso.
La disminución de peso es una técnica de regularización que penaliza los valores grandes de los parámetros añadiendo un término proporcional a la suma de los cuadrados de los parámetros a la función de pérdida. En SGD estándar, la disminución de peso es equivalente a la regularización L2, pero esta equivalencia se rompe en métodos adaptativos como Adam. AdamW fue diseñado para restaurar el comportamiento previsto de la disminución de peso en optimizadores adaptativos.
El optimizador Adam
Adam (Estimación Adaptativa de Momentos) fue introducido por Diederik Kingma y Jimmy Ba en 2014. Mantiene tasas de aprendizaje por parámetro mediante el mantenimiento de un promedio móvil exponencial de los gradientes pasados (primer momento) y de los gradientes pasados al cuadrado (segundo momento). La regla de actualización para Adam es:
\[ \theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \]
donde \(\hat{m}_t\) y \(\hat{v}_t\) son estimaciones corregidas por sesgo del primer y segundo momento, \(\eta\) es la tasa de aprendizaje y \(\epsilon\) es una pequeña constante para la estabilidad numérica. Adam se hizo popular debido a su robustez a los hiperparámetros y su rápida convergencia, especialmente para entrenar redes profundas.
Sin embargo, en la implementación original de Adam, la disminución de peso se implementaba como regularización L2, que añade un término \(\frac{\lambda}{2} \|\theta\|^2\) a la pérdida. Este término se incluye entonces en el gradiente, y debido a que Adam normaliza el gradiente por el segundo momento, la disminución de peso efectiva se convierte en \(\lambda / \sqrt{\hat{v}_t}\). Esto significa que los parámetros con gradientes grandes reciben menos regularización, lo que puede conducir a un sobreajuste y a una mala generalización.
El algoritmo AdamW
AdamW modifica la regla de actualización eliminando el término de regularización L2 del gradiente y aplicando en su lugar la disminución de peso directamente a los parámetros después de la actualización adaptativa. La regla de actualización se convierte en:
\[ \theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} - \eta \lambda \theta_t \]
donde \(\lambda\) es el coeficiente de disminución de peso. Este desacoplamiento asegura que la disminución de peso se aplique uniformemente a todos los parámetros, independientemente de sus magnitudes de gradiente. Los autores argumentaron que esto conduce a una mejor generalización y a un entrenamiento más estable, especialmente cuando se utilizan tasas de aprendizaje grandes.
En su artículo, Loshchilov y Hutter demostraron que AdamW supera a Adam con regularización L2 en varias tareas de referencia, incluida la clasificación de imágenes en CIFAR-10 y el modelado de lenguaje en Penn Treebank. También mostraron que AdamW es más robusto a la elección de la tasa de aprendizaje y del coeficiente de disminución de peso.
Impacto en el aprendizaje profundo
AdamW ha tenido un impacto significativo en el campo de la inteligencia artificial. Ahora es el optimizador predeterminado en muchos marcos populares de aprendizaje profundo, como PyTorch y TensorFlow, y se utiliza ampliamente en el entrenamiento de transformers y grandes modelos de lenguaje. Por ejemplo, muchos modelos desarrollados por organizaciones como OpenAI, Anthropic y Google DeepMind utilizan AdamW como parte de sus pipelines de entrenamiento.
La disminución de peso desacoplada en AdamW ha sido particularmente beneficiosa para entrenar modelos a gran escala, donde la regularización es crucial para prevenir el sobreajuste en conjuntos de datos masivos. También se ha demostrado que mejora la velocidad de convergencia y el rendimiento final en comparación con otros optimizadores como SGD con momento o Adam con regularización L2.
Comparación con otros optimizadores
AdamW se compara a menudo con otros optimizadores adaptativos como AdaGrad, RMSProp y Adam. Mientras que AdaGrad y RMSProp ajustan las tasas de aprendizaje basándose en gradientes históricos, Adam combina tanto el momento como las tasas de aprendizaje adaptativas. AdamW mejora a Adam al corregir el problema de la disminución de peso, haciéndolo más fiable para una amplia gama de tareas.
Otro optimizador relacionado es SGD con momento, que es más simple pero a menudo requiere un ajuste cuidadoso del programa de tasas de aprendizaje. AdamW ofrece un buen equilibrio entre facilidad de uso y rendimiento, razón por la cual se ha convertido en la opción preferida para muchos profesionales. Sin embargo, algunos estudios han demostrado que SGD con momento puede lograr una mejor generalización en ciertas tareas si se ajusta adecuadamente, pero AdamW sigue siendo competitivo y más robusto a las elecciones de hiperparámetros.
Consideraciones prácticas
Al usar AdamW, hay algunas consideraciones prácticas. El coeficiente de disminución de peso \(\lambda\) se establece típicamente en un valor pequeño, como 0.01 o 0.1, pero puede necesitar ajustes para tareas específicas. La tasa de aprendizaje a menudo se establece en un valor como 1e-4 o 3e-4 para entrenar transformers. Además, AdamW a menudo se beneficia de un programa de calentamiento de la tasa de aprendizaje, donde la tasa de aprendizaje aumenta gradualmente desde un valor pequeño hasta el valor objetivo durante los primeros miles de pasos.
En la práctica, AdamW está implementado en la mayoría de las bibliotecas de aprendizaje profundo, por lo que los usuarios pueden simplemente especificar el optimizador y establecer el parámetro de disminución de peso. Por ejemplo, en PyTorch, se puede usar torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01). Esta simplicidad ha contribuido a su adopción generalizada.
Extensiones y variantes
Se han propuesto varias extensiones y variantes de AdamW. Por ejemplo, AdamW con disminución de peso desacoplada se ha combinado con otras técnicas como Lookahead, que mantiene un conjunto más lento de parámetros para actualizaciones más suaves. Otra variante es AdamW con programas de tasa de aprendizaje de recocido de coseno, que ha demostrado mejorar el rendimiento en tareas de clasificación de imágenes.
En el contexto de los grandes modelos de lenguaje, AdamW se utiliza a menudo con entrenamiento de precisión mixta y acumulación de gradientes para manejar tamaños de lote grandes. Algunos marcos también ofrecen implementaciones fusionadas de AdamW que reducen el uso de memoria y mejoran la eficiencia computacional, lo cual es importante al entrenar modelos con miles de millones de parámetros.
Conclusión
AdamW se ha convertido en una herramienta fundamental en el kit de herramientas del aprendizaje automático. Al desacoplar la disminución de peso de las actualizaciones adaptativas de gradiente, aborda un defecto sutil pero importante en el optimizador Adam original, lo que conduce a una mejor generalización y a un entrenamiento más estable. Su simplicidad y eficacia lo han convertido en la opción predeterminada para muchas aplicaciones de aprendizaje profundo, desde la clasificación de imágenes hasta el procesamiento del lenguaje natural. A medida que el campo continúa evolucionando, AdamW sigue siendo un método de optimización fiable y ampliamente utilizado.
Referencias
- Loshchilov, I., & Hutter, F. (2017). Decoupled Weight Decay Regularization. arXiv:1711.05101.
- Kingma, D. P., & Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXiv:1412.6980.