O Surpreendente Acordo entre a Teoria da Otimização Convexa e o Agendamento de Taxa de Aprendizado para Treinamento de Modelos Grandes

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

January 31, 2025

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

cs.AI

Resumo

Demonstramos que os cronogramas de taxa de aprendizado para o treinamento de modelos grandes se comportam de forma surpreendentemente semelhante a um limite de desempenho da teoria da otimização convexa não suave. Fornecemos um limite para o cronograma constante com redução linear; em particular, o benefício prático da redução é refletido no limite devido à ausência de termos logarítmicos. Além disso, demonstramos que essa correspondência surpreendentemente próxima entre teoria e prática de otimização pode ser explorada para ajuste da taxa de aprendizado: alcançamos melhorias perceptíveis no treinamento de modelos do tipo Llama de 124M e 210M através de (i) estender o cronograma para treinamento contínuo com taxa de aprendizado ótima e (ii) transferir a taxa de aprendizado ótima entre cronogramas.

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.