MARS: Sfruttare il Potere della Riduzione della Varianza per l'Addestramento di Modelli Grandi

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

November 15, 2024

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

cs.AI

Abstract

L'addestramento delle reti neurali profonde - e più recentemente, dei modelli di grandi dimensioni - richiede ottimizzatori efficienti e scalabili. Gli algoritmi di gradiente adattivi come Adam, AdamW e le loro varianti sono stati centrali per questo compito. Nonostante lo sviluppo di numerosi algoritmi di riduzione della varianza nell'ultimo decennio mirati ad accelerare l'ottimizzazione stocastica sia in contesti convessi che non convessi, la riduzione della varianza non ha trovato un ampio successo nell'addestramento delle reti neurali profonde o dei grandi modelli linguistici. Di conseguenza, è rimasto un approccio meno preferito nell'IA moderna. In questo articolo, per sbloccare il potenziale della riduzione della varianza per un addestramento efficiente dei grandi modelli, proponiamo un framework di ottimizzazione unificato, MARS (Make vAriance Reduction Shine), che concilia i metodi di gradiente precondizionati con la riduzione della varianza tramite una tecnica di momento stocastico scalato ricorsivo. All'interno del nostro framework, introduciamo tre istanze di MARS che sfruttano gli aggiornamenti del gradiente precondizionato basati su AdamW, Lion e Shampoo, rispettivamente. Stabiliamo anche un collegamento tra i nostri algoritmi e gli ottimizzatori esistenti. I risultati sperimentali sull'addestramento dei modelli GPT-2 indicano che MARS supera costantemente AdamW di gran lunga.

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.