MARS: Liberando el Poder de la Reducción de Varianza para Entrenar 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

Resumen

Entrenar redes neuronales profundas, y más recientemente, modelos grandes, requiere optimizadores eficientes y escalables. Algoritmos de gradiente adaptativos como Adam, AdamW y sus variantes han sido fundamentales para esta tarea. A pesar del desarrollo de numerosos algoritmos de reducción de varianza en la última década destinados a acelerar la optimización estocástica en entornos convexas y no convexas, la reducción de varianza no ha tenido un éxito generalizado en el entrenamiento de redes neuronales profundas o modelos de lenguaje grandes. En consecuencia, ha seguido siendo un enfoque menos preferido en la IA moderna. En este documento, para liberar el poder de la reducción de varianza para el entrenamiento eficiente de modelos grandes, proponemos un marco de optimización unificado, MARS (Make vAriance Reduction Shine), que concilia los métodos de gradiente precondicionado con la reducción de varianza a través de una técnica de momento estocástico escalado recursivo. Dentro de nuestro marco, presentamos tres instancias de MARS que aprovechan actualizaciones de gradiente precondicionado basadas en AdamW, Lion y Shampoo, respectivamente. También establecemos una conexión entre nuestros algoritmos y optimizadores existentes. Los resultados experimentales en el entrenamiento de modelos GPT-2 indican que MARS supera consistentemente a AdamW por un amplio margen.

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.