Summary
The paper focuses on learning rate schedules in scaling large language models (LLMs). Traditionally, LLMs are trained with a cosine learning rate schedule, which requires a pre-specified training duration for learning rate decay. This makes it difficult to dynamically adjust the training duration, as early stopping and continued training after the decay result in suboptimal model performance. To address this, the authors study an alternative schedule of constant learning rate with cooldown, which enables one to get the optimal model at any step during the training process. The authors systematically compare different choices of cool down duration and cool down schedule, and demonstrate that constant learning rate with 1 - sqrt cooldown schedule matches the performance of cosine schedule tuned for that duration.
Beyond cooldown, the authors compare two different approaches of removing the need of learning rate schedule: stochastic weight averaging and schedule-free optimizer. The authors show that while they exhibit promising performance, the performance of models obtained in the middle of the training still cannot match the performance of the proposed constant learning rate with cooldown.
Finally, the authors validate the proposed learning rate schedule in the scaling law study of compute optimal models, and demonstrate that the constant learning rate with cooldown schedule leads to the same results as cosine learning rate schedule, while saving half of the compute flops.
Strengths
The paper presents a simple approach to address a major limitation of cosine learning rate schedule, and demonstrates its effectiveness with solid experimentations. Overall I think the paper is of very high quality.
The proposed method is simple and highly practical, which makes it relevant for a broad range of today’s scaling studies. It is widely known that scaling is the key to the success of LLMs, and it is important to establish compute optimal scaling laws as large scaling training is really expensive. The proposed learning rate schedule is able to halve the flops of the chinchilla compute optimal horizon study without adding extra implementation or tuning complexity, which I believe is a major contribution.
The proposed method is supported with solid experimentations. The authors include comprehensive comparisons of different schedules of cooldown and hyperparameter choices, which convincingly establish the advantage of the proposed schedule and can also serve as tuning guide when adopting it. The authors also compare the proposed method to popular alternative approaches.
The paper is well written. The main idea and the experiments are laid out in a logical way that makes it easy to follow. The experiment results are also presented in an intuitive way, making it easy to compare methods.
Weaknesses
While the experiment results are very solid, I do hope that the authors could expand the scale of experiments to larger models. Larger models are usually more sensitive to suboptimal hyperparameters, and therefore it would be great if the proposed method can also be verified in larger models, such as a 3B or 7B model which are common sizes for practical LLMs.
Questions
Another highly relevant use of training LLMs without fixed duration is continued pre-training, which is quite common in practice as people often continue the pre-training on specialized data for different use cases (such as turning general models into code specialized models). It would be great if the authors could include some discussions and experiments on how the proposed constant learning rate with cooldown can be applied in continued pre-training. For example, one question to ask would be would it be better to continue from the checkpoint before or after cooldown.
Limitations
The authors have sufficiently addressed the limitations and societal impacts.