Plus d'Adam : La mise à l'échelle du taux d'apprentissage à l'initialisation est tout ce dont vous avez besoin

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

December 16, 2024

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

cs.AI

Résumé

Dans ce travail, nous remettons en question la nécessité des méthodes de gradient adaptatif pour l'entraînement des réseaux de neurones profonds. SGD-SaI est une amélioration simple mais efficace de la descente de gradient stochastique avec momentum (SGDM). SGD-SaI effectue une Mise à l'échelle du taux d'apprentissage à l'initialisation (SaI) pour des groupes de paramètres distincts, guidée par leurs rapports signal sur bruit de gradient respectifs (g-SNR). En ajustant les taux d'apprentissage sans recourir à un moment d'ordre deux adaptatif, SGD-SaI aide à prévenir les déséquilibres d'entraînement dès la première itération et réduit de moitié l'utilisation de la mémoire de l'optimiseur par rapport à AdamW. Malgré sa simplicité et son efficacité, SGD-SaI égale ou surpasse de manière constante AdamW dans l'entraînement de diverses tâches basées sur des Transformers, surmontant efficacement un défi de longue date lié à l'utilisation de SGD pour l'entraînement de Transformers. SGD-SaI excelle dans la classification ImageNet-1K avec les Vision Transformers (ViT) et la pré-entraînement GPT-2 pour de grands modèles de langage (LLMs, uniquement des décodeurs de transformer), démontrant sa robustesse face aux variations d'hyperparamètres et sa praticité pour des applications diverses. Nous avons également testé sa robustesse sur des tâches telles que l'adaptation fine LoRA pour les LLMs et les modèles de diffusion, où il surpasse de manière constante les optimiseurs de pointe. Du point de vue de l'efficacité mémoire, SGD-SaI permet d'économiser de manière significative la mémoire des états de l'optimiseur, réduisant l'utilisation de la mémoire de 5,93 Go pour GPT-2 (1,5 milliard de paramètres) et de 25,15 Go pour Llama2-7B par rapport à AdamW dans des paramètres d'entraînement en pleine précision.

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.