Geen Adam meer: Schalen van de leersnelheid bij initialisatie is alles wat je nodig hebt.

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

Samenvatting

In dit werk stellen we de noodzaak van adaptieve gradiëntmethoden voor het trainen van diepe neurale netwerken ter discussie. SGD-SaI is een eenvoudige maar effectieve verbetering van stochastische gradiëntdaling met momentum (SGDM). SGD-SaI voert Learning Rate Scaling at Initialization (SaI) uit naar verschillende parametergroepen, geleid door hun respectievelijke gradiënt signaal-ruisverhoudingen (g-SNR). Door leerpercentages aan te passen zonder te vertrouwen op adaptieve tweede-orde momentum, helpt SGD-SaI om trainingsonevenwichtigheden vanaf de allereerste iteratie te voorkomen en vermindert het het geheugengebruik van de optimizer met de helft in vergelijking met AdamW. Ondanks zijn eenvoud en efficiëntie, presteert SGD-SaI consequent even goed of beter dan AdamW bij het trainen van verschillende op Transformer gebaseerde taken, waarbij het effectief een langdurige uitdaging overwint van het gebruik van SGD voor het trainen van Transformers. SGD-SaI blinkt uit in ImageNet-1K classificatie met Vision Transformers (ViT) en GPT-2 pretraining voor grote taalmodellen (LLM's, alleen transformer-decoder), waarbij het robuustheid toont tegen hyperparameter variaties en praktisch nut voor diverse toepassingen. We hebben ook de robuustheid ervan getest op taken zoals LoRA fine-tuning voor LLM's en diffusiemodellen, waar het consequent beter presteert dan state-of-the-art optimalisatoren. Vanuit een geheugenefficiëntieperspectief behaalt SGD-SaI aanzienlijke geheugenbesparingen voor optimizerstaten, waarbij het geheugengebruik met 5,93 GB vermindert voor GPT-2 (1,5B parameters) en 25,15 GB voor Llama2-7B in vergelijking met AdamW in instellingen voor training met volledige precisie.

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.