Addio ad Adam: la regolazione del tasso di apprendimento all'inizializzazione è tutto ciò di cui hai bisogno

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

December 16, 2024

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

cs.AI

Abstract

In questo lavoro, mettiamo in discussione la necessità dei metodi di gradiente adattivi per l'addestramento delle reti neurali profonde. SGD-SaI è un potenziamento semplice ma efficace della discesa del gradiente stocastico con momento (SGDM). SGD-SaI esegue la Scalatura del tasso di apprendimento all'Inizializzazione (SaI) per gruppi di parametri distinti, guidati dai rispettivi rapporti segnale-rumore del gradiente (g-SNR). Regolando i tassi di apprendimento senza fare affidamento sul momento adattivo del secondo ordine, SGD-SaI aiuta a prevenire squilibri nell'addestramento fin dalla prima iterazione e riduce l'utilizzo della memoria dell'ottimizzatore della metà rispetto ad AdamW. Nonostante la sua semplicità ed efficienza, SGD-SaI si allinea costantemente o supera AdamW nell'addestramento di una varietà di compiti basati su Transformer, superando efficacemente una sfida di lunga data nell'uso di SGD per l'addestramento dei Transformer. SGD-SaI eccelle nella classificazione ImageNet-1K con Vision Transformers (ViT) e nel preaddestramento GPT-2 per modelli linguistici di grandi dimensioni (LLM, solo decoder transformer), dimostrando robustezza alle variazioni degli iperparametri e praticità per applicazioni diverse. Abbiamo inoltre testato la sua robustezza in compiti come il fine-tuning LoRA per LLM e modelli di diffusione, dove supera costantemente ottimizzatori all'avanguardia. Dal punto di vista dell'efficienza della memoria, SGD-SaI ottiene risparmi di memoria sostanziali per gli stati dell'ottimizzatore, riducendo l'utilizzo della memoria di 5,93 GB per GPT-2 (1,5 miliardi di parametri) e di 25,15 GB per Llama2-7B rispetto ad AdamW nelle impostazioni di addestramento a precisione 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.