Não Mais Adam: Escalonamento da Taxa de Aprendizado na Inicialização é Tudo o que Você Precisa

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

December 16, 2024

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

cs.AI

Resumo

Neste trabalho, questionamos a necessidade de métodos de gradiente adaptativo para o treinamento de redes neurais profundas. O SGD-SaI é uma melhoria simples, porém eficaz, do descida de gradiente estocástica com momento (SGDM). O SGD-SaI realiza o Dimensionamento da Taxa de Aprendizado na Inicialização (SaI) para grupos de parâmetros distintos, guiados por suas respectivas razões sinal-ruído do gradiente (g-SNR). Ao ajustar as taxas de aprendizado sem depender de momento adaptativo de segunda ordem, o SGD-SaI ajuda a evitar desequilíbrios no treinamento desde a primeira iteração e reduz o uso de memória do otimizador pela metade em comparação com o AdamW. Apesar de sua simplicidade e eficiência, o SGD-SaI consistentemente corresponde ou supera o desempenho do AdamW no treinamento de uma variedade de tarefas baseadas em Transformadores, superando efetivamente um desafio de longa data de usar SGD para treinar Transformadores. O SGD-SaI se destaca na classificação do ImageNet-1K com Transformadores de Visão (ViT) e no pré-treinamento do GPT-2 para grandes modelos de linguagem (LLMs, apenas decodificador de transformador), demonstrando robustez a variações de hiperparâmetros e praticidade para aplicações diversas. Testamos ainda sua robustez em tarefas como ajuste fino LoRA para LLMs e modelos de difusão, onde consistentemente supera otimizadores de ponta. Do ponto de vista da eficiência de memória, o SGD-SaI alcança economias substanciais de memória para estados do otimizador, reduzindo o uso de memória em 5,93 GB para o GPT-2 (1,5 bilhão de parâmetros) e 25,15 GB para o Llama2-7B em comparação com o AdamW em configurações de treinamento em precisão total.

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.