Technical Summary: Predictable Scaling Laws of Optimal Hyperparameters for LLM Continued Pre-training

Problem Statement

Continued Pre-training (CPT) is a critical mechanism for adapting Large Language Models (LLMs) to domain-specific applications, offering a compute-efficient alternative to training from scratch. However, the efficacy of CPT is highly sensitive to hyperparameter configurations, particularly learning rate (LR) and batch size (B). Current practices rely heavily on heuristics or exhaustive grid searches, which are computationally expensive, prone to training instability, and lack principled guidance.

Two fundamental challenges hinder the application of scaling laws to CPT:

1. Existence and Characterization: It remains unverified whether predictable, monotonic relationships exist between compute budget and optimal hyperparameters when training does not start from random initialization.
2. Quantifying Initial State: Unlike scratch training, CPT begins from a non-zero optimization state. The "black box" nature of a pre-trained checkpoint's progress on a new, shifted data distribution makes it difficult to quantify its effective contribution to the training process.

Methodology

The authors propose a two-stage, model-agnostic framework to establish quantitative relationships between compute budget and optimal hyperparameters, enabling zero-shot prediction for any given checkpoint.

Stage 1: Empirical Law Discovery

The authors first characterize the scaling behavior by training small-scale proxy models (100M and 500M parameters) across diverse hyperparameter configurations on the target CPT data.

• Iso-Loss Curves: They construct curves in the (Batch Size, Learning Rate) plane to identify the minimum compute required to reach a specific validation loss. The vertices of these curves define the optimal hyperparameter pairs $(B_{opt}, LR_{opt})$ for a given loss level.
• Functional Fitting: They observe that optimal hyperparameters follow power-law trends relative to validation loss. They fit functions $B_{opt} = f(L)$ and $LR_{opt} = g(L)$.
• Compute-to-Hyperparameter Mapping: By combining these with standard loss-compute scaling laws ($L(C) = L_0 + \alpha C^{-\gamma}$), they derive direct mappings from total compute budget $C$ to optimal hyperparameters:
  $$B_{opt}(C) = f(L_0 + \alpha \cdot C^{-\gamma})$$
  $$LR_{opt}(C) = g(L_0 + \alpha \cdot C^{-\gamma})$$

Stage 2: State-Aware Hyperparameter Prediction

To apply these laws to arbitrary checkpoints, the authors introduce the concept of Equivalent Pre-training Compute ($C_{pre}$).

• Definition: $C_{pre}$ represents the theoretical compute budget required to reach a checkpoint's current validation loss if it had been trained from random initialization on the target domain.
• Calculation: By inverting the loss-compute scaling law, $C_{pre}$ is calculated from the checkpoint's initial validation loss ($L_{init}$):
  $$C_{pre} = \left( \frac{\alpha}{L_{init} - L_0} \right)^{1/\gamma}$$
• Total Effective Compute: The total compute budget for the CPT run is defined as the sum of the inherited state and the planned budget: $C_{total} = C_{pre} + C_{cpt}$.
• Prediction: The optimal hyperparameters for the target run are predicted by evaluating the derived scaling functions at $C_{total}$.

Key Contributions

1. Scaling Laws for CPT: The paper empirically establishes that optimal hyperparameters in continued pre-training follow stable, predictable scaling laws. Specifically, optimal batch size increases monotonically with compute budget, while optimal learning rate decreases.
2. Equivalent Pre-training Compute ($C_{pre}$): A novel metric is proposed to formally quantify the training state inherited from initial checkpoints. This transforms an arbitrary starting point into a measurable coordinate on a continuous optimization trajectory, addressing the non-zero initialization challenge.
3. Efficient Framework: The authors validate a framework that predicts optimal hyperparameters without trial runs on target models, reducing hyperparameter search overhead by up to 90% compared to traditional grid searches.

Experimental Results

The framework was evaluated on Dense-8B and MoE-3B models using a 55B-token corpus spanning general knowledge, mathematics, and code.

• Performance: Models trained with predicted hyperparameters achieved peak average scores (51.3 for Dense-8B and 57.9 for MoE-3B) across diverse benchmarks (MMLU, GSM8K, HumanEval, etc.), outperforming off-target settings.
• Ablation Studies:
  • Learning Rate: Excessively high LRs caused performance collapse (catastrophic forgetting), while low LRs constrained knowledge acquisition efficiency. The predicted LR balanced these factors.
  • Batch Size: Small batches led to unstable optimization, while excessively large batches reduced gradient update frequency, hindering convergence. The predicted batch size optimized this trade-off.
• Efficiency: Compared to grid searches on target models, the proposed method reduced computational costs by 92.0% for Dense-8B and 73.1% for MoE-3B by leveraging small-scale proxy models.
• Validity of $C_{pre}$: Ablation studies showed that ignoring the initial state (Variant A) or using raw pre-training token counts (Variant B) resulted in suboptimal hyperparameters and lower performance, confirming that $C_{pre}$ effectively calibrates for domain shifts.

Significance and Claims

The paper claims to provide a principled and efficient methodology for diverse continued pre-training scenarios. By demonstrating that optimal hyperparameters follow predictable scaling laws even when starting from non-random initialization, the work offers a viable alternative to costly heuristics and grid searches. The introduction of $C_{pre}$ allows practitioners to treat CPT as a continuation of a single optimization path, enabling reliable hyperparameter prediction for models at any stage of their lifecycle. The authors note that while the framework generalizes across architectures, verification on ultra-large scales (70B+ parameters) and highly specialized domains remains a direction for future investigation.