CMR Scaling Law: Predicting Critical Mixture Ratios for Continual Pre-training of Language Models (2024)

Abstract: Large Language Models (LLMs) excel in diverse tasks but often underperform in specialized fields due to limited domain-specific or proprietary corpus. Continual pre-training (CPT) enhances LLM capabilities by imbuing new domain-specific or proprietary knowledge while replaying general corpus to prevent catastrophic forgetting. The data mixture ratio of general corpus and domain-specific corpus, however, has been chosen heuristically, leading to sub-optimal training efficiency in practice. In this context, we attempt to re-visit the scaling behavior of LLMs under the hood of CPT, and discover a power-law relationship between loss, mixture ratio, and training tokens scale. We formalize the trade-off between general and domain-specific capabilities, leading to a well-defined Critical Mixture Ratio (CMR) of general and domain data. By striking the balance, CMR maintains the model’s general ability and achieves the desired domain transfer, ensuring the highest utilization of available resources. Considering the balance between efficiency and effectiveness, CMR can be regarded as the optimal mixture ratio. Through extensive experiments, we ascertain the predictability of CMR, propose CMR scaling law and have substantiated its generalization. These findings offer practical guidelines for optimizing LLM training in specialized domains, ensuring both general and domain-specific performance while efficiently managing training resources.

## introduction large language models (llms) exhibit versatile abilities, including question answering, translation, summarization, role-playing, etc. @xcite @xcite . their performance, however, may degrade in specific domains due to limited corre- † equal contribution * corresponding author sponding pre-training data. to enhance llms' abilities in specialized areas and avoid the enormous cost of re-training, a popular approach is continual pre-training (cpt) @xcite @xcite . this approaches are likely to equip llms with new domain-related capabilities without much general performance penalty. although cpt has been proven effective on multiple domains such as code @xcite , law @xcite and medicine @xcite , the interplay among loss prediction and its scaling behavior with model size, and the number of training tokens is yet to be fully explored. additionally, the composition of continual pre-training data is simply set up in a heuristic manner @xcite , far from being principled. an inappropriate mixture ratio can lead to inefficient training (requiring excessive computation to adapt to specific domains) or insufficient training (failing to adequately reduce domain-specific loss). in light of this, three question hurdles we need to cross are as follows: does the optimal data mixture ratio exist for cpt? if so, how does it evolve with model scale or training token volume? are there any involved simple yet principled laws? currently, several studies examine the scaling laws associated with different data mixture ratios. for instance, @xcite investigate how data mixtures shape scaling laws in the pre-training phase from the ground up, while @xcite seek to pinpoint the optimal data mixture ratio in cpt, but overlook its crucial connection with the essential trade-off between general and domain loss in cpt. therefore, to strengthen our understanding about cpt and guide the experiments in the future, we attempt to address these questions with empirical studies on cpt of llms. specifically, we pre-train several llms with different model sizes from scratch and perform cpt on downstream domains (finance and academic papers) with different datamixture ratios. our main contributions can be summarized as follows: formalization of the trade-off in cpt. we formalize the balance between domain-specific and general abilities during cpt by introducing the concept of feasible mixture ratios. cpt under feasible mixture ratios maintains performance on general data while enhancing performance on domain-specific data. we identify the maximum feasible mixture ratio as the critical mixture ratio (cmr), and regard it as the optimal mixture ratio by our definition. ## key results we train a series of llms with multiple mixture ratios of domain-specific data and general data to analyse the scaling behaviour in cpt. the method is detailed in § 3.2. based on our experimental setup, we summarize the key results as follows: ## background and methods the scaling law in the pre-training stage has been widely studied. in this work, we simplify the form of scaling law as much as possible, which is essentially consistent with previous works in § 6. in this section, we will elaborate on the three main concepts involved in this work, including objective of cpt (definition 1), feasible mixture ratio (definition 2) and cmr (definition 3). then, we describe our experiment setups, including data preparation, experiment procedures and evaluation. ## does the critical mixture ratio exist? a larger mixture ratio implies a higher proportion of domain-specific data in the training set, resulting in a lower domain loss. however, due to the potential catastrophic forgetting of domain transfer, it is essential to ensure that the loss in the new domain continues to decrease while the original capabilities of llms are preserved and not compromised during cpt. consequently, a higher mixture ratio is not always best. this raises an important question: does a critical mixture ratio (cmr) exist that can balance these two goals of cpt in definition 1 effectively and efficiently? figure 1 (left) demonstrates that for models of various sizes, there is at least one curve at a specific ratio that shows a downward trend, highlighted by yellow dotted arrows. this indicates the presence of feasible mixture ratios that align with our cpt objective. on the other hand, larger models tend to have bigger feasible mixture ratios set f (more curves with yellow dotted arrows). for curves that meet the objective of cpt, a higher ratio is preferable, as it incorporates more domain knowledge while optimizing training efficiency within the tolerance of decline in general capacity. therefore, the critical mixture ratio is defined as the highest proportion among these satisfactory curves, representing the optimal ratio for the given model size and limited training token volume. if feasible ratios exist, we can conclude that cmr is also supposed to exist. fundamentally, the existence of cmr arises from the trade-off between general and domain-specific capabilities, as well as the limited data and computing resources. according to definition 3, the cmr is present across models of different scales, as shown in figure 1 . this figure illustrates the existence of cmr as the maximum value within the feasible set. however, the precise value of cmr can not be determined from the figure, as it requires extensive experiments with different mixture ratios. the estimation of cmrs is discussed in 5.3 and plotted in figure 5 . to look closely, we enlarged the longitudinal section of m 940m in the 3d graph in figure 1 and placed it on the right side. it can be seen that as the mixture ratio increases, the curve continues to rise until the loss in the general domain exceeds our tolerance. the potentially controversial issue is that the downward trend in one-third of the curves is as clear as in the rest. the reason why it is feasible curve here can be found in appendix b.2. although it is not easily noticeable, there are indeed points on this curve where the slope is less than 0. findings from another perspective, we plot the loss curves of the models under the same mixture ratio as shown in figure 2 . when the mixture ratio is 1/4, all models can achieve the training objective of cpt. however, at a 1/3 mixture ratio, only m 940m , m 1.6b and m 3.1b achieve the cpt goal. this indicates that cmr for m 940m is around 1/4 within the scope of our training token volumes, while the cmrs for m 940m , m 1.6b and m 3.1b are at least 1/3. in other words, cmrs slightly increase with model size, suggesting that larger models can accommodate a higher proportion of domain data. we also further this finding by taking more cross-sections of figure 1 (left) in appendix e and the predicted cmr in following § 5. this phenomenon can be explained by the models' ability to consume domain knowledge. as the proportion of domain-specific data increases, the knowledge that the model needs to learn also increases. llms with smaller size struggle to absorb much of domain knowledge while preserving the general knowledge, leading to a degradation in their original general performance. in contrast, models with larger sizes can accommodate more knowledge with more parameters, thereby maintaining better performance. the existence of cmr indicates that in the process of cpt, we may explore the cmr scaling law to seek the best mixture ratio under resource constraints and domain data limitations, thereby optimizing training effectiveness and efficiency. in other words, the next question to answer is whether we can predict the cmr for model m s given a maximum amount of continuation training token volume, t max . to this end, two basic prerequisites must be met: predicting losses for different mixture ratio and predicting losses for different training token volume. in this section, we will demonstrate that these two prerequisites have been satisfied separately in § 5.1 and § 5.2, and finally detail the scaling law to predict cmr in § 5.3. to keep notations simple, we omit fixed variables in the loss function (l dom/gen and l ∆dom/∆gen ) in this following. ## conclusion in this work, we investigated the scaling behavior of llms under continual pre-training (cpt) to address the limitations of domain-specific performance. we provided a clear definition of critical mixture ratio (cmr) for optimizing the mixture ratio of general and domain-specific data. our experiments revealed a power-law relationship between loss, mixture ratio, and training data size, allowing us to predict the cmr efficiently. these findings may offer practical guidelines for optimizing llm training, helping to balance general and domain-specific performance while minimizing resource consumption. additionally, our study suggests that understanding the cpt process and scaling laws could be valuable for future research aimed at enhancing llm capabilities in specialized fields.