MARS: Die Kraft der Varianzreduktion für das Training großer Modelle entfesseln

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

November 15, 2024

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

cs.AI

Zusammenfassung

Das Training von tiefen neuronalen Netzwerken - und in jüngster Zeit auch großer Modelle - erfordert effiziente und skalierbare Optimierer. Adaptive Gradienten-Algorithmen wie Adam, AdamW und deren Varianten waren dabei entscheidend. Trotz der Entwicklung zahlreicher Algorithmen zur Varianzreduktion im vergangenen Jahrzehnt, die darauf abzielten, stochastische Optimierung in konvexen und nicht-konvexen Umgebungen zu beschleunigen, hat die Varianzreduktion im Training von tiefen neuronalen Netzwerken oder großen Sprachmodellen keine weitreichenden Erfolge erzielt. Folglich ist sie ein weniger bevorzugter Ansatz in der modernen KI geblieben. In diesem Paper schlagen wir vor, die Kraft der Varianzreduktion für das effiziente Training großer Modelle freizusetzen, indem wir ein vereinheitlichtes Optimierungsframework namens MARS (Make vAriance Reduction Shine) vorstellen, das vorkonditionierte Gradientenmethoden mit Varianzreduktion über eine skalierte stochastische rekursive Momentum-Technik in Einklang bringt. Innerhalb unseres Frameworks stellen wir drei Instanzen von MARS vor, die auf vorkonditionierten Gradientenupdates basieren, die auf AdamW, Lion und Shampoo basieren. Wir ziehen auch eine Verbindung zwischen unseren Algorithmen und bestehenden Optimierern. Experimentelle Ergebnisse zum Training von GPT-2-Modellen zeigen, dass MARS AdamW deutlich übertrifft.

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.