L'étonnante concordance entre la théorie de l'optimisation convexe et la planification du taux d'apprentissage pour l'entraînement de grands modèles

The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training

January 31, 2025

Auteurs: Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, Francis Bach

cs.AI

Résumé

Nous montrons que les programmes de taux d'apprentissage pour l'entraînement de grands modèles se comportent de manière surprenamment similaire à une limite de performance de la théorie de l'optimisation convexe non lisse. Nous fournissons une limite pour le programme constant avec refroidissement linéaire ; en particulier, le bénéfice pratique du refroidissement se reflète dans la limite en raison de l'absence de termes logarithmiques. De plus, nous montrons que cette correspondance étonnamment proche entre la théorie de l'optimisation et la pratique peut être exploitée pour l'ajustement du taux d'apprentissage : nous obtenons des améliorations notables pour l'entraînement de modèles de type Llama de 124M et 210M en (i) étendant le programme pour un entraînement continu avec un taux d'apprentissage optimal, et (ii) en transférant le taux d'apprentissage optimal entre les programmes.

English

We show that learning-rate schedules for large model training behave surprisingly similar to a performance bound from non-smooth convex optimization theory. We provide a bound for the constant schedule with linear cooldown; in particular, the practical benefit of cooldown is reflected in the bound due to the absence of logarithmic terms. Further, we show that this surprisingly close match between optimization theory and practice can be exploited for learning-rate tuning: we achieve noticeable improvements for training 124M and 210M Llama-type models by (i) extending the schedule for continued training with optimal learning-rate, and (ii) transferring the optimal learning-rate across schedules.