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

TL;DR

This paper introduces SGD-SaI, a simple learning rate scaling method at initialization that replaces adaptive optimizers like AdamW, achieving comparable or better performance in training Transformers while reducing memory usage.

Contribution

The authors propose SGD-SaI, a novel learning rate scaling technique at initialization that eliminates the need for adaptive gradient methods in training deep neural networks.

Findings

SGD-SaI matches or outperforms AdamW in Transformer training tasks.

SGD-SaI significantly reduces memory usage compared to AdamW.

SGD-SaI demonstrates robustness across various tasks and hyperparameters.

Abstract

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.