MARS : Libérer la Puissance de la Réduction de la Variance pour l'Entraînement de Grands Modèles

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

Résumé

L'entraînement des réseaux neuronaux profonds – et plus récemment, des grands modèles – nécessite des optimiseurs efficaces et évolutifs. Les algorithmes de gradient adaptatifs tels qu'Adam, AdamW et leurs variantes ont été essentiels à cette tâche. Malgré le développement de nombreux algorithmes de réduction de la variance au cours de la dernière décennie visant à accélérer l'optimisation stochastique dans des contextes convexes et non convexes, la réduction de la variance n'a pas connu un succès généralisé dans l'entraînement des réseaux neuronaux profonds ou des grands modèles de langage. Par conséquent, elle est restée une approche moins privilégiée dans l'IA moderne. Dans cet article, pour libérer la puissance de la réduction de la variance pour un entraînement efficace des grands modèles, nous proposons un cadre d'optimisation unifié, MARS (Make vAriance Reduction Shine), qui concilie les méthodes de gradient préconditionnées avec la réduction de la variance via une technique d'élan stochastique récursive mise à l'échelle. Au sein de notre cadre, nous introduisons trois exemples de MARS qui exploitent des mises à jour de gradient préconditionnées basées sur AdamW, Lion et Shampoo, respectivement. Nous établissons également un lien entre nos algorithmes et les optimiseurs existants. Les résultats expérimentaux sur l'entraînement des modèles GPT-2 indiquent que MARS surpasse systématiquement AdamW de manière significative.

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.