MARS: Het Benutten van de Kracht van Variantiereductie voor het Trainen van Grote Modellen

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

November 15, 2024

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

cs.AI

Samenvatting

Het trainen van diepe neurale netwerken - en meer recentelijk, grote modellen - vereist efficiënte en schaalbare optimalisatoren. Adaptieve gradiëntalgoritmen zoals Adam, AdamW en hun varianten zijn cruciaal gebleken voor deze taak. Ondanks de ontwikkeling van talrijke algoritmen voor variantiereductie in het afgelopen decennium, gericht op het versnellen van stochastische optimalisatie in zowel convexe als niet-convexe omgevingen, heeft variantiereductie geen brede acceptatie gevonden bij het trainen van diepe neurale netwerken of grote taalmodellen. Als gevolg hiervan is het een minder populaire benadering gebleven in de moderne AI. In dit artikel stellen we voor om de kracht van variantiereductie voor efficiënte training van grote modellen te benutten door middel van een verenigd optimalisatiekader, MARS (Make vAriance Reduction Shine), dat voorwaardelijke gradiëntmethoden verzoent met variantiereductie via een geschaalde stochastische recursieve momentumtechniek. Binnen ons kader introduceren we drie instanties van MARS die gebruikmaken van voorwaardelijke gradiëntupdates op basis van AdamW, Lion en Shampoo, respectievelijk. We leggen ook een verband tussen onze algoritmen en bestaande optimalisatoren. Experimentele resultaten bij het trainen van GPT-2 modellen tonen aan dat MARS consequent AdamW aanzienlijk overtreft.

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.