Learning Dynamics in Continual Pre-Training for Large Language Models

Learning Dynamics in Continual Pre-Training for Large Language Models

Speakers: Xingjin Wang, University of Chinese Academy of Sciences; Howe Tissue; Lu Wang; Linjing Li; Daniel Zeng

Conference: ICML 2025

YouTube: https://slideslive.com/39043895

Overview

This talk, presented by Xingjin Wang from the University of Chinese Academy of Sciences, delves into the intricate learning dynamics of continual pre-training (CPT) for Large Language Models (LLMs). Continual pre-training is a critical technique that enables LLMs to rapidly adapt to new, specialized domains such as mathematics or code, extend their context windows, and integrate new knowledge for subsequent downstream tasks like reinforcement learning (RL) fine-tuning. Despite its widespread use, the CPT process is characterized by numerous interacting variables, including learning rate schedules, data replay ratios, and initialization checkpoints, making its optimization a complex challenge.

The core of this research addresses two fundamental questions: first, can a comprehensive framework account for all variables influencing CPT's final performance? And second, can performance be accurately predicted throughout the entire CPT process, rather than merely at its conclusion? The speakers introduce a novel CPT scaling law that extends their prior work on pre-training dynamics, offering a unified mathematical framework to model the evolution of loss during continual pre-training. This law not only provides a powerful tool for understanding the underlying mechanisms of adaptation but also offers practical guidance for optimizing CPT strategies and developing more effective open-source models.

The significance of this work lies in its potential to transform how LLMs are continually updated and deployed. By demystifying the complex interplay of hyperparameters and data distributions in CPT, the proposed scaling law empowers researchers and practitioners to predict model behavior, optimize training resources, and enhance the adaptability and performance of LLMs across diverse applications. It provides a principled approach to navigate the trade-offs inherent in CPT, moving beyond heuristic trial-and-error to data-driven optimization.

Background

▶ Watch: Introduction to continual pre-training and research questions (0:00)

The foundation of this research builds upon the speakers' prior work on pre-training learning dynamics, specifically a scaling law with learning rate annealing. This foundational work posited a scaling law function capable of fitting and predicting the full loss curve during pre-training, not just the final loss point. They conceptualized the loss surface as a "cylinder," comprising two principal directions: a forward direction and an annealing direction.

When a model is trained with a constant learning rate, the loss typically decreases smoothly along the forward direction, representing the primary progression of learning. However, once the learning rate begins to decay – a common practice in modern deep learning to stabilize training and improve convergence – the loss exhibits a rapid decrease along the annealing direction. This phenomenon is crucial for understanding how models fine-tune their internal representations. To quantify these dynamics, the prior work introduced concepts like the forward area (the integral area below the learning rate scheduler curve) and the annealing area (the integral area above the learning rate decay rate from the scheduler). These areas provide a mathematical description of how the learning rate schedule influences the loss dynamics at each training step.

The challenge in extending this understanding to continual pre-training (CPT) is multifaceted. CPT inherently involves transitioning a pre-trained model from one data distribution (the original pre-training dataset, or PT dataset) to a new, often specialized, data distribution (the continual pre-training dataset, or CPT dataset). This transition introduces a distribution shift that complicates the learning dynamics. Furthermore, CPT processes often involve numerous hyperparameters, such as the peak learning rate for the CPT phase, the data replay ratio (the proportion of original PT data mixed with new CPT data), and the specific learning rate scheduler employed. The goal is to understand how these variables interact to affect the model's performance on both the original and new domains, as well as its overall generalization capabilities. Prior to this work, a holistic, predictive framework for these complex dynamics in CPT was largely absent, leaving practitioners to rely on extensive empirical tuning.

Key Findings

▶ Watch: Extending to CPT: CPT loss curves and distribution shift (2:00)

The research presents several pivotal findings that collectively offer a comprehensive understanding and predictive framework for continual pre-training:

1. A Unified CPT Scaling Law: The most significant contribution is the development of a CPT scaling law equation that accurately models and predicts the entire loss curve during continual pre-training. This law successfully integrates the effects of learning rate annealing (from the original pre-training dynamics) with a newly characterized distribution shift term. This term quantifies the discrepancy arising from transitioning between different data distributions.

1. Characterization of Distribution Shift: The distribution shift is found to exhibit a consistent pattern, largely independent of the specific transfer starting point or checkpoint. It can be effectively modeled using a power law form. This discovery simplifies the modeling process, allowing for robust prediction even when the exact starting state of the CPT is varied.

1. The Concept of Loss Potential: The authors introduce loss potential as a novel metric, defined as the "height" of the current model state on the loss surface in the annealing view, representing its potential for future loss reduction via learning rate annealing. A key finding is that pre-trained models exhibiting a higher loss potential consistently achieve lower CPT validation loss. This suggests a crucial implication: when releasing open-source models, it is more beneficial to release versions with higher loss potential to facilitate more effective and efficient subsequent continual pre-training by users.

1. Hyperparameter Optimization Capabilities: The CPT scaling law enables a detailed analysis of key hyperparameters:

• Replay Ratio: A larger replay ratio (mixing more original pre-training data with new CPT data) leads to a smaller distribution shift, indicating better retention of original knowledge.
• Peak Learning Rate: The scaling law accurately predicts the entire loss curve for different peak learning rates. It reveals a typical trade-off: a larger peak learning rate often leads to a higher initial pre-training validation loss but can result in a lower continual pre-training validation loss, highlighting the need for careful balancing based on specific objectives.
• Optimal Hyperparameter Calculation: The framework can calculate optimal values for loss potential, peak learning rate, and replay ratio by assigning normalized coefficients to the pre-training and continual pre-training validation losses, allowing for objective-driven hyperparameter selection.

1. Adaptability to Black-Box Models: The CPT scaling law remains applicable even when dealing with "black-box" open-source pre-trained models where detailed pre-training information is unavailable. This is achieved by utilizing a proxy PT dataset to approximate the original distribution and tracing specific parameters required for fitting the equation. This significantly expands the practical utility of the framework.

1. RL Learning Dynamics Insights: Beyond CPT for LLMs, the research also extends to Reinforcement Learning (RL) learning dynamics. It identifies that while stable entropy contributes to stable RL training, annealing entropy can accelerate local performance improvement and enable faster convergence. This led to the proposal of entropy scheduling, analogous to learning rate scheduling, allowing for effective control of entropy within a single training run to maximize reward.

Technical Deep Dive

▶ Watch: Key Factor Analysis: Loss potential and its benefits (4:00)

The technical core of this work lies in extending the established pre-training (PT) scaling law to encompass the complexities of continual pre-training (CPT). The original PT scaling law modeled loss dynamics based on forward area and annealing area, derived from the learning rate scheduler. For CPT, this concept is adapted by replacing the forward and annealing areas with the sum of those areas accumulated across both the pre-training and continual pre-training phases. This allows the model to account for the cumulative effect of learning rate schedules over the entire training continuum.

A critical aspect of CPT is the distribution shift that occurs when a model trained on one dataset (PT data) is exposed to a new dataset (CPT data). To isolate and characterize this shift, the authors conducted experiments where models were pre-trained and continually pre-trained using the same constant learning rate. This strategic choice eliminated the confounding effects of learning rate annealing, allowing them to observe the distribution shift term in isolation. They found that this distribution shift term tended to overlap regardless of the specific transfer starting point or checkpoint. This striking consistency enabled them to model the distribution shift using a power law form, simplifying its integration into the overall scaling law.

The final CPT scaling law equation is a powerful formulation that combines two primary effects: the influence of learning rate annealing (derived from the extended forward and annealing areas) and the newly characterized distribution shift term. This comprehensive equation allows for the accurate fitting and prediction of the entire CPT loss curve. For experimental validation, the researchers utilized WSD (Warmup-Stable-Decay) and cosine learning rate schedulers for both pre-training and continual pre-training the models.

To further dissect the CPT process, the authors introduced hidden PT curves.

1. The first is a "hidden PT curve trained on PT dataset" (blue dashed line), representing a hypothetical scenario where the model continues training only on the original PT dataset with the same learning rate scheduler as the CPT phase.
2. The second is a "hidden PT curve trained on CPT dataset" (orange dashed line), representing a model trained from scratch on the CPT dataset with identical training setups.

The actual CPT loss curve (solid orange line) is observed to initially deviate from the blue dashed line (due to the distribution shift) and then converge towards the orange dashed line (as it adapts to the CPT data distribution). The discrepancy between the blue dashed line and the actual CPT loss curve is precisely what they term the distribution shift.

The concept of loss potential is introduced by visualizing the loss surface as a transition between two states. In this "annealing view," the height of the current model state is defined as its loss potential, signifying its inherent capacity for future loss reduction through learning rate annealing. A steeper slope on this transition surface indicates a larger distance between the PT and CPT data distributions. The finding that higher loss potential leads to lower CPT validation loss is crucial for model release strategies.

The CPT scaling law also provides a robust framework for hyperparameter analysis. For instance, it predicts that a larger replay ratio will result in a smaller distribution shift, implying that retaining more of the original data helps mitigate catastrophic forgetting and smooth the transition. Similarly, the equation can predict the entire loss curve for different peak learning rates, illustrating the trade-off between pre-training validation loss and continual pre-training validation loss. To optimize for specific objectives, the framework allows for the assignment of normalized coefficients to these two validation losses. By calculating the optimal loss potential, peak learning rate, and replay ratio based on these coefficients, practitioners can tailor training to their desired outcomes. For out-of-domain validation sets, a linear combination of pre-training and continual pre-training validation losses can be used for evaluation.

Addressing the practical challenge of black-box open-source models, where detailed pre-training logs are unavailable, the authors propose a clever workaround. They suggest using a proxy PT dataset to approximate the original distribution of the pre-training data. By tracing specific parameters that require fitting through a short CPT run, their CPT scaling law can still accurately fit and predict the loss curve of models like Llama 32 (1 billion parameters), demonstrating its broad applicability.

Finally, the talk briefly touches upon RL learning dynamics, focusing on entropy. They observe that while stable entropy ensures training stability, annealing entropy can significantly accelerate local performance improvements and lead to faster convergence. This led to the proposal of entropy scheduling, which combines stable and annealing phases, analogous to learning rate scheduling, to effectively control entropy within a single training run and rapidly boost rewards.

Experimental Setup & Results

▶ Watch: Analyzing variables: Replay ratio, peak LR, and optimization (4:40)

The experimental validation of the proposed CPT scaling law and its underlying concepts was conducted using a Llama 32 model with 1 billion parameters. This choice of a moderately sized LLM allowed for thorough experimentation while providing results generalizable to larger models.

Key aspects of the experimental setup included:

• Learning Rate Schedulers: The models were pre-trained and continually pre-trained using both WSD (Warmup-Stable-Decay) and cosine learning rate schedulers. These are common and effective schedulers in LLM training, providing a robust testbed for the scaling law's ability to model diverse annealing behaviors.
• Training Steps: The models were trained for "long enough tokens," specifically mentioning 16,000 steps, ensuring that the learning dynamics had sufficient time to manifest and stabilize, avoiding premature convergence issues.
• Validation Metrics: Performance was assessed using:
• Pre-training validation loss (PT validation loss): To track the model's performance on the original domain.
• Continual pre-training validation loss (CPT validation loss): To track adaptation to the new domain.
• Out-of-domain validation set: To evaluate generalization beyond the specific CPT dataset.

The experimental results provided strong empirical support for the theoretical framework:

1. Accurate Loss Curve Prediction: The CPT scaling law equation demonstrated remarkable accuracy in fitting and predicting the entire CPT loss curve. This was a central claim and was visually supported by figures showing the predicted curve closely aligning with the actual training loss. This ability to predict the full curve, rather than just the final point, represents a significant advancement.

1. Impact of Loss Potential: Experiments confirmed that pre-trained models with a higher loss potential consistently achieved a lower CPT validation loss. This finding underscores the importance of the model's "readiness" for further adaptation and provides a concrete metric for evaluating the suitability of pre-trained checkpoints for CPT.

1. Replay Ratio Effects: A clear relationship was observed between the replay ratio and the distribution shift. Specifically, a larger replay ratio (i.e., including more original PT data during CPT) led to a smaller distribution shift. This quantitatively validates the intuition that replaying old data helps mitigate catastrophic forgetting and smooth the transition to new domains.

1. Peak Learning Rate Trade-offs: The CPT scaling law successfully predicted the loss curves for different peak learning rates during CPT. The results illustrated a common trade-off: a larger peak learning rate often led to a higher PT validation loss (suggesting some degradation on the original domain) but a lower CPT validation loss (faster and better adaptation to the new domain). This highlights the need for a principled approach to selecting the peak learning rate based on the specific objectives of the CPT task.

1. Optimal Hyperparameter Calculation: The framework's ability to calculate optimal values for loss potential, peak learning rate, and replay ratio was demonstrated. By allowing users to assign normalized coefficients to PT and CPT validation losses, the system could identify the hyperparameter settings that best balanced these competing objectives for a given training goal.

1. Black-Box Model Applicability: The method for adapting the CPT scaling law to black-box open-source models was successfully demonstrated using a Llama 32 1B model. By employing a proxy PT dataset and fitting parameters from a short initial CPT run, the law was still able to accurately fit and predict the subsequent loss curve, proving its practical utility in real-world scenarios where full pre-training details are unavailable.

1. RL Entropy Scheduling: The exploration into RL learning dynamics revealed that entropy scheduling, particularly the annealing of entropy, led to rapid increases in reward and model performance within short training steps. This finding provides an analogous mechanism to learning rate scheduling for optimizing RL agents.

The consistency of these results across various experimental conditions and the alignment with the proposed theoretical framework strongly validate the utility and predictive power of the CPT scaling law.

Practical Implications

▶ Watch: Beyond CPT: RL learning dynamics and entropy scheduling (6:30)

The research on learning dynamics in continual pre-training carries significant practical implications for various stakeholders in the AI/ML ecosystem, from model developers to deployment teams.

For Model Builders and Open-Source Model Releasers:

The introduction of loss potential as a key metric offers a novel criterion for evaluating and releasing pre-trained models. Instead of simply releasing the model with the lowest final pre-training loss, this work suggests that releasing a version with higher loss potential is more beneficial. Such models are inherently better primed for subsequent continual pre-training, leading to lower CPT validation loss and faster adaptation for downstream users. This could lead to a shift in how open-source model repositories curate and present their offerings, potentially including a "CPT-ready" version alongside a standard one.

For Practitioners and Infra Teams Optimizing LLMs:

The CPT scaling law provides a powerful analytical tool that can significantly reduce the extensive trial-and-error typically involved in optimizing continual pre-training.

• Hyperparameter Optimization: Practitioners can use the law to predict the impact of different peak learning rates and replay ratios on both original domain performance and new domain adaptation. This enables a more principled approach to selecting hyperparameters, moving from empirical guesses to data-driven decisions. The ability to calculate optimal hyperparameters for specific objectives means less compute spent on hyperparameter sweeps.
• Predictive Power: The capacity to predict the entire CPT loss curve from a few initial steps is invaluable. This allows teams to forecast long-term training outcomes, identify potential issues early, and make informed decisions about stopping criteria or adjusting strategies without waiting for full training runs.
• Resource Efficiency: By optimizing hyperparameters and predicting outcomes, organizations can significantly save computational resources (GPU/TPU hours) and time, accelerating the development and deployment cycles of specialized LLMs.

For Deployers and Adaptation Teams:

The framework's applicability to black-box open-source models is a crucial practical advantage. Many organizations leverage pre-trained models without access to their full training history. The ability to approximate the original data distribution using a proxy PT dataset and fit the CPT scaling law parameters means that even with limited information, these teams can still leverage the predictive power of the framework to optimize CPT for their specific use cases. This democratizes advanced CPT optimization techniques.

Trade-offs and Limitations:

The research highlights inherent trade-offs, such as the relationship between a larger peak learning rate leading to a higher PT validation loss but a lower CPT validation loss. Practitioners must carefully weigh these trade-offs based on their specific application. If retaining strong performance on the original domain is paramount, a lower peak learning rate might be preferred, even if it means slower adaptation to the new domain. Conversely, if rapid and deep adaptation to the new domain is the goal, a higher peak learning rate might be acceptable.

While powerful, the model relies on fitting parameters from observed data. The robustness of these fitted parameters across vastly different model architectures, scales, or data distributions (beyond those tested) would need further investigation. The "power law form" for distribution shift, while effective, might have limitations in extreme cases of distribution divergence. The talk also briefly mentions RL learning dynamics and entropy scheduling, suggesting that similar scaling law principles could be applied to optimize RL fine-tuning processes, offering faster reward accumulation and convergence for agents.

In essence, this work provides a robust analytical lens for understanding and controlling the complex learning dynamics of continual pre-training, transforming it from an art into a more precise science.

Key Takeaways

• CPT Dynamics Are Predictable: A novel CPT scaling law accurately models and predicts the entire loss curve during continual pre-training by combining the effects of learning rate annealing and distribution shift.
• Distribution Shift is Quantifiable: The discrepancy caused by switching data distributions in CPT can be consistently modeled using a power law form, simplifying its integration into the scaling law.
• "Loss Potential" is a Key Metric: Pre-trained models with higher loss potential (their inherent capacity for future loss reduction) consistently achieve lower CPT validation loss, suggesting a new criterion for releasing CPT-friendly models.
• Hyperparameter Optimization is Possible: The CPT scaling law enables the analysis and calculation of optimal replay ratios and peak learning rates, allowing practitioners to tailor CPT strategies for specific objectives (e.g., balancing original vs. new domain performance).
• Black-Box Models Can Be Optimized: Even for open-source models without full training details, the CPT scaling law can be applied using a proxy PT dataset and parameter fitting from short CPT runs, significantly expanding its practical utility.
• Entropy Scheduling for RL: Analogous to learning rate scheduling, annealing entropy can significantly accelerate local performance improvement and convergence in Reinforcement Learning, offering a new optimization lever for RL fine-tuning.

About the Speaker(s)

The primary presenter of this work was Xingjin Wang, representing the University of Chinese Academy of Sciences. He introduced the research and detailed the methodologies and findings. The co-authors contributing to this work include Howe Tissue, Lu Wang, Linjing Li, and Daniel Zeng. While specific titles and affiliations beyond Xingjin Wang's institution were not detailed in the transcript, their collective contribution to this technical article on learning dynamics in continual pre-training for large language models highlights a collaborative effort from researchers involved in advanced machine learning and language model development.