Die überraschende Übereinstimmung zwischen der Theorie der konvexen Optimierung und der Lernratenplanung für das Training großer Modelle

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

January 31, 2025

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

cs.AI

Zusammenfassung

Wir zeigen, dass Lernratenpläne für das Training großer Modelle überraschend ähnlich zu einer Leistungsgrenze aus der Theorie der nicht-glatten konvexen Optimierung verhalten. Wir liefern eine Grenze für den konstanten Plan mit linearem Abkühlen; insbesondere spiegelt sich der praktische Nutzen des Abkühlens in der Grenze aufgrund des Fehlens von logarithmischen Termen wider. Weiterhin zeigen wir, dass diese überraschend enge Übereinstimmung zwischen Optimierungstheorie und Praxis für die Feinabstimmung der Lernrate genutzt werden kann: Wir erzielen deutliche Verbesserungen beim Training von 124M und 210M Llama-Modellen, indem wir (i) den Plan für das fortgesetzte Training mit optimaler Lernrate erweitern und (ii) die optimale Lernrate zwischen Plänen übertragen.

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.