Kein Adam mehr: Skalierung der Lernrate bei der Initialisierung ist alles, was Sie brauchen.

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

December 16, 2024

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

cs.AI

Zusammenfassung

In dieser Arbeit hinterfragen wir die Notwendigkeit adaptiver Gradientenmethoden für das Training tiefer neuronaler Netzwerke. SGD-SaI ist eine einfache, aber effektive Verbesserung des stochastischen Gradientenabstiegs mit Momentum (SGDM). SGD-SaI führt das Lernraten-Scaling bei der Initialisierung (SaI) für verschiedene Parametergruppen durch, geleitet durch ihre jeweiligen Gradienten-Signal-Rausch-Verhältnisse (g-SNR). Durch die Anpassung der Lernraten, ohne auf adaptives Momentum zweiter Ordnung angewiesen zu sein, hilft SGD-SaI, Trainingsungleichgewichte bereits ab der ersten Iteration zu verhindern und den Speicherverbrauch des Optimierers im Vergleich zu AdamW zu halbieren. Trotz seiner Einfachheit und Effizienz übertrifft SGD-SaI konstant oder erreicht die Leistung von AdamW beim Training verschiedener Aufgaben auf Basis von Transformern, wodurch eine langjährige Herausforderung bei der Verwendung von SGD für das Training von Transformern effektiv bewältigt wird. SGD-SaI zeichnet sich bei der ImageNet-1K-Klassifizierung mit Vision-Transformern (ViT) und dem GPT-2-Feintuning für große Sprachmodelle (LLMs, nur Transformer-Decoder) aus, was seine Robustheit gegenüber Hyperparameter-Variationen und seine Praktikabilität für vielfältige Anwendungen zeigt. Wir haben seine Robustheit auch bei Aufgaben wie LoRA-Feintuning für LLMs und Diffusionsmodellen getestet, bei denen er konsequent bessere Leistungen als optimale Optimierer erzielt. Aus der Perspektive der Speichereffizienz erzielt SGD-SaI erhebliche Speichereinsparungen für Optimiererzustände und reduziert den Speicherverbrauch um 5,93 GB für GPT-2 (1,5 Mrd. Parameter) und 25,15 GB für Llama2-7B im Vergleich zu AdamW in Einstellungen für das Training in voller Präzision.

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.