L'Inaspettato Accordo Tra la Teoria dell' Ottimizzazione Convessa e la Pianificazione del Tasso di Apprendimento per l'Addestramento di Modelli di Grandi Dimensioni

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

January 31, 2025

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

cs.AI

Abstract

Dimostriamo che gli schemi di tasso di apprendimento per l'addestramento di modelli di grandi dimensioni si comportano in modo sorprendentemente simile a un limite di prestazione della teoria dell'ottimizzazione convessa non liscia. Forniamo un limite per lo schema costante con raffreddamento lineare; in particolare, il beneficio pratico del raffreddamento è riflesso nel limite a causa dell'assenza di termini logaritmici. Inoltre, dimostriamo che questa corrispondenza sorprendentemente stretta tra teoria e pratica dell'ottimizzazione può essere sfruttata per l'ottimizzazione del tasso di apprendimento: otteniamo miglioramenti significativi per l'addestramento di modelli di tipo Llama da 124M e 210M mediante (i) l'estensione dello schema per l'addestramento continuato con tasso di apprendimento ottimale e (ii) il trasferimento del tasso di apprendimento ottimale tra gli schemi.

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.