No más Adam: Escalar la tasa de aprendizaje en la inicialización es todo lo que necesitas.

No More Adam: Learning Rate Scaling at Initialization is All You Need

December 16, 2024

Autores: Minghao Xu, Lichuan Xiang, Xu Cai, Hongkai Wen

cs.AI

Resumen

En este trabajo, cuestionamos la necesidad de los métodos de gradiente adaptativo para el entrenamiento de redes neuronales profundas. SGD-SaI es una mejora simple pero efectiva de la descenso de gradiente estocástico con momento (SGDM). SGD-SaI realiza Escalado de tasa de aprendizaje en la Inicialización (SaI) para grupos de parámetros distintos, guiados por sus respectivas relaciones señal-ruido de gradiente (g-SNR). Al ajustar las tasas de aprendizaje sin depender de un momento adaptativo de segundo orden, SGD-SaI ayuda a prevenir desequilibrios en el entrenamiento desde la primera iteración y reduce a la mitad el uso de memoria del optimizador en comparación con AdamW. A pesar de su simplicidad y eficiencia, SGD-SaI coincide o supera consistentemente a AdamW en el entrenamiento de una variedad de tareas basadas en Transformadores, superando efectivamente un desafío de larga data en el uso de SGD para entrenar Transformadores. SGD-SaI destaca en la clasificación de ImageNet-1K con Transformadores de Visión (ViT) y en el preentrenamiento de GPT-2 para modelos de lenguaje grandes (LLMs, solo decodificador de transformador), demostrando robustez ante variaciones de hiperparámetros y practicidad para diversas aplicaciones. Probamos además su robustez en tareas como el ajuste fino de LoRA para LLMs y modelos de difusión, donde supera consistentemente a optimizadores de última generación. Desde una perspectiva de eficiencia de memoria, SGD-SaI logra ahorros sustanciales de memoria para los estados del optimizador, reduciendo el uso de memoria en 5.93 GB para GPT-2 (1.5B parámetros) y 25.15 GB para Llama2-7B en comparación con AdamW en configuraciones de entrenamiento de precisión completa.

English

In this work, we question the necessity of adaptive gradient methods for training deep neural networks. SGD-SaI is a simple yet effective enhancement to stochastic gradient descent with momentum (SGDM). SGD-SaI performs learning rate Scaling at Initialization (SaI) to distinct parameter groups, guided by their respective gradient signal-to-noise ratios (g-SNR). By adjusting learning rates without relying on adaptive second-order momentum, SGD-SaI helps prevent training imbalances from the very first iteration and cuts the optimizer's memory usage by half compared to AdamW. Despite its simplicity and efficiency, SGD-SaI consistently matches or outperforms AdamW in training a variety of Transformer-based tasks, effectively overcoming a long-standing challenge of using SGD for training Transformers. SGD-SaI excels in ImageNet-1K classification with Vision Transformers(ViT) and GPT-2 pretraining for large language models (LLMs, transformer decoder-only), demonstrating robustness to hyperparameter variations and practicality for diverse applications. We further tested its robustness on tasks like LoRA fine-tuning for LLMs and diffusion models, where it consistently outperforms state-of-the-art optimizers. From a memory efficiency perspective, SGD-SaI achieves substantial memory savings for optimizer states, reducing memory usage by 5.93 GB for GPT-2 (1.5B parameters) and 25.15 GB for Llama2-7B compared to AdamW in full-precision training settings.