A Learning Rate Path Switching Training Paradigm for Version Updates of Large Language Models (2024)

Abstract: Due to the continuous emergence of new data, version updates have become an indispensable requirement for Large Language Models (LLMs). The training paradigms for version updates of LLMs include pre-training from scratch (PTFS) and continual pre-training (CPT). Preliminary experiments demonstrate that PTFS achieves better pre-training performance, while CPT has lower training cost. Moreover, their performance and training cost gaps widen progressively with version updates. To investigate the underlying reasons for this phenomenon, we analyze the effect of learning rate adjustments during the two stages of CPT: preparing an initialization checkpoint and continual pre-training based on this checkpoint. We find that a large learning rate in the first stage and a complete learning rate decay process in the second stage are crucial for version updates of LLMs. Hence, we propose a learning rate path switching training paradigm. Our paradigm comprises one main path, where we pre-train a LLM with the maximal learning rate, and multiple branching paths, each of which corresponds to an update of the LLM with newly-added training data. Extensive experiments demonstrate the effectiveness and generalization of our paradigm. Particularly, when training four versions of LLMs, our paradigm reduces the total training cost to 58% compared to PTFS, while maintaining comparable pre-training performance.

## introduction in recent years, there has been significant progress in the research of large language models (llms). by performing large-scale training on massive datasets, llms have demonstrated remarkable capabilities, contributing to various fields @xcite work was done when zhihao wang, shiyu liu and jianheng huang were interning at vivo ai lab. in this paper, we firstly conduct preliminary experiments to compare the above two paradigms in version updates of llms. compared with ptfs, cpt uses previous checkpoints for initialization, resulting in lower total training cost. however, cpt suffers from inferior pre-training performance, which becomes increasingly serious as version updates progress. to study the reasons for this phenomenon, we break down the cpt process into two stages: the first stage involves preparing an initialization checkpoint, and the second stage performing continual pre-training based on this checkpoint. then, we conduct two groups of experiments to analyze the effect of learning rate adjustments during these two stages, leading to two conclusions: 1) the larger the learning rate in the first stage, the better the performance of updated llms in the second stage; 2) for the second stage, a complete learning rate decay process is beneficial to ensure the optimal performance of updated llms. based on the above analyses, we propose a learning rate path switching training paradigm for version updates of llms. to better illustrate our paradigm, we take the most commonly used cosine learning rate schedule (smith and topin, 2019) as an example, and plot the learning rate curves of ptfs, cpt and our paradigm in figure 1 . please note that our paradigm is also applicable to other schedules, such as knee @xcite , and multi-step @xcite learning rate schedules. in short, the learning rate curve of our paradigm comprises one main path and multiple branching paths, each of which corresponds to a version update of llm. as shown by the main path in figure 1 , we pre-train a llm with the maximal learning rate, providing superior initialization checkpoints for subsequent continual pre-training. when we want to update the llm with newly-added training data, we perform continual pre-training on the llm with a dynamically-adjusted learning rate. referring back to figure 1 , after a few steps of training with the maximal learning rate, the learning rate fast decays to its minimum, which effectively ensures the training performance of the updated llm. meanwhile, on the main path, we continue to pre-train the original checkpoint with the maximal learning rate, facilitating subsequent llm updates. our paradigm better balances model performance and training cost compared to the other two paradigms, as detailed in figure 2 . to summarize, our main contributions are as follows: ## preliminary study in this section, we first compare the performance of ptfs and cpt in version updates of llms, and then analyze the underlying reasons for their performance gap. ## our paradigm based on the conclusions from section 2, we propose a learning rate path switching paradigm for version updates of llms in this section. the training cost of our paradigm is lower than that of ptfs, and it achieves significantly better performance than cpt, with performance even comparable to that of ptfs. ## experiment in this section, we still use the settings of the preliminary study to conduct more experiments, comparing the performance and training cost of different training paradigms. ## conclusion and future work this paper focuses on how to effectively balance model performance and training cost for version updates of llms. we begin by comparing two training paradigms: ptfs and cpt, concluding that ptfs achieves better pre-training performance, while cpt has lower training cost. through the analysis in the preliminary study, we find that 1) a large learning rate is beneficial for providing better initialization checkpoints for subsequent updates, and 2) a complete learning rate decay process enables the updated llms to achieve optimal performance. based on the above two findings, we propose a learning rate path switching paradigm for version updates of llms, which comprises one main path and multiple branching paths. on the main path, we pre-train the llms with the maximal learning rate to provide superior initialization checkpoints for subsequent updates. when an update is required, our paradigm switches from the main path to a branching path, undergoing a complete learning rate decay process. experimental results and further analyses strongly demonstrate the effectiveness and generalization of our paradigm. in the future, we will further expand the practical scope of our paradigm. current research mainly focuses on the pre-training phase and does not include supervised fine-tuning, safety alignment, etc., which could be integrated into the fast-decaying stage of our paradigm. additionally, we plan to explore the applicability of our paradigm in the context of multimodal large language models.