MARS: Liberando o Poder da Redução de Variância para Treinar Modelos Grandes

MARS: Unleashing the Power of Variance Reduction for Training Large Models

November 15, 2024

Autores: Huizhuo Yuan, Yifeng Liu, Shuang Wu, Xun Zhou, Quanquan Gu

cs.AI

Resumo

O treinamento de redes neurais profundas - e mais recentemente, modelos grandes - exige otimizadores eficientes e escaláveis. Algoritmos de gradiente adaptativos como Adam, AdamW e suas variantes têm sido fundamentais para essa tarefa. Apesar do desenvolvimento de inúmeros algoritmos de redução de variância na última década com o objetivo de acelerar a otimização estocástica em ambientes convexos e não convexos, a redução de variância não obteve amplo sucesso no treinamento de redes neurais profundas ou grandes modelos de linguagem. Consequentemente, tem sido uma abordagem menos favorecida na IA moderna. Neste artigo, para liberar o potencial da redução de variância para o treinamento eficiente de grandes modelos, propomos um framework de otimização unificado, MARS (Make vAriance Reduction Shine), que concilia métodos de gradiente precondicionados com redução de variância por meio de uma técnica de momento estocástico escalonado recursivo. Dentro do nosso framework, introduzimos três instâncias de MARS que aproveitam atualizações de gradiente precondicionadas com base em AdamW, Lion e Shampoo, respectivamente. Também estabelecemos uma conexão entre nossos algoritmos e otimizadores existentes. Resultados experimentais no treinamento de modelos GPT-2 indicam que o MARS supera consistentemente o AdamW por uma grande margem.

English

Training deep neural networks--and more recently, large models--demands efficient and scalable optimizers. Adaptive gradient algorithms like Adam, AdamW, and their variants have been central to this task. Despite the development of numerous variance reduction algorithms in the past decade aimed at accelerating stochastic optimization in both convex and nonconvex settings, variance reduction has not found widespread success in training deep neural networks or large language models. Consequently, it has remained a less favored approach in modern AI. In this paper, to unleash the power of variance reduction for efficient training of large models, we propose a unified optimization framework, MARS (Make vAriance Reduction Shine), which reconciles preconditioned gradient methods with variance reduction via a scaled stochastic recursive momentum technique. Within our framework, we introduce three instances of MARS that leverage preconditioned gradient updates based on AdamW, Lion, and Shampoo, respectively. We also draw a connection between our algorithms and existing optimizers. Experimental results on training GPT-2 models indicate that MARS consistently outperforms AdamW by a large margin.