대규모 모델 훈련을 위한 볼록 최적화 이론과 학습 속도 일정 간의 놀라운 일치

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

January 31, 2025

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

cs.AI

초록

대규모 모델 훈련을 위한 학습률 스케줄이 부드럽지 않은 볼록 최적화 이론의 성능 한계와 놀랍도록 유사하게 작동하는 것을 보여줍니다. 우리는 선형 쿨다운을 가진 상수 스케줄에 대한 한계를 제시합니다. 특히, 쿨다운의 실용적 이점은 로그항이 없어서 한계에 반영됩니다. 더 나아가 최적화 이론과 실무 간 놀랍도록 밀접한 일치가 학습률 조정에 활용될 수 있음을 보여줍니다: 우리는 최적 학습률로 지속적 훈련을 위한 스케줄을 확장하고, 최적 학습률을 스케줄 간 전이함으로써 124M 및 210M Llama 유형 모델의 훈련에서 현저한 개선을 달성합니다.

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.