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

TL;DR

This paper introduces the CMR scaling law, a power-law relationship that predicts the optimal mixture ratio of general and domain-specific data in continual pre-training of LLMs, enhancing efficiency and performance.

Contribution

It formalizes the Critical Mixture Ratio (CMR) and demonstrates its ability to optimize data mixture for improved LLM continual pre-training.

Findings

Discovered a power-law relationship between loss, mixture ratio, and training scale.

Defined the Critical Mixture Ratio (CMR) for balancing general and domain-specific capabilities.

Validated the CMR scaling law through extensive experiments.

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.