MARS: Разблокирование потенциала снижения дисперсии для обучения больших моделей

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

November 15, 2024

Авторы: Huizhuo Yuan, Yifeng Liu, Shuang Wu, Xun Zhou, Quanquan Gu

cs.AI

Аннотация

Обучение глубоких нейронных сетей - и в последнее время больших моделей - требует эффективных и масштабируемых оптимизаторов. Адаптивные алгоритмы градиентного спуска, такие как Adam, AdamW и их вариации, стали ключевыми в этой задаче. Несмотря на разработку множества алгоритмов снижения дисперсии за последнее десятилетие, направленных на ускорение стохастической оптимизации как в выпуклых, так и невыпуклых сценариях, снижение дисперсии не получило широкого признания при обучении глубоких нейронных сетей или крупных языковых моделей. Следовательно, это остается менее предпочтительным подходом в современном искусственном интеллекте. В данной статье мы предлагаем унифицированную оптимизационную структуру, MARS (Make vAriance Reduction Shine), чтобы раскрыть потенциал снижения дисперсии для эффективного обучения крупных моделей, которая совмещает методы предварительной обработки градиента с снижением дисперсии с помощью масштабируемой стохастической рекурсивной техники импульса. В рамках нашей структуры мы представляем три варианта MARS, которые используют обновления предварительно обработанного градиента на основе AdamW, Lion и Shampoo соответственно. Мы также устанавливаем связь между нашими алгоритмами и существующими оптимизаторами. Экспериментальные результаты по обучению моделей GPT-2 показывают, что MARS постоянно превосходит AdamW на значительное расстояние.

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.