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

Abstract: The efficacy of continued pre-training for LLMs hinges upon hyperparameter configurations, such as learning rate and batch size. However, current practices often rely on heuristics or grid searches, leading to training instability and excessive costs. In this work, we first empirically discover that optimal hyperparameters follow stable and predictable scaling laws throughout the continued pre-training process. Leveraging these insights, we propose a novel framework to establish quantitative relationships between compute budget and optimal hyperparameters for a given checkpoint. Our approach has two stages: (1) \textit{Empirical Law Discovery}, where we train small-scale proxy models to derive functions mapping compute budget to optimal hyperparameters via standard loss-compute scaling laws; and (2) \textit{State-Aware Hyperparameter Prediction}, where we evaluate an initial checkpoint's validation loss and use the inverse scaling law to estimate its \textit{equivalent pre-training compute} -- the compute needed to achieve the same loss from scratch. Combining this with the planned compute budget, we predict optimal hyperparameters for the target run. Empirical results demonstrate that our method reduces the hyperparameter search overhead by up to 90\% while achieving comparable or superior performance relative to baselines. This model-agnostic framework generalizes across architectures, providing a principled and efficient methodology for diverse continued pre-training scenarios starting from any given point.

• The paper introduces a two-stage framework that uses compute-based scaling laws and equivalent pre-training compute to predict optimal learning rates and batch sizes for new LLM checkpoints without target-model grid searches.
• Experiments show that optimal batch size increases while learning rate decreases with compute, with architecture-specific patterns indicating that MoE models require substantially greater batch scaling than dense models.
• The predicted settings matched or exceeded grid-search configurations on Dense-8B and MoE-3B models, reducing hyperparameter search compute by 92.0% and 73.1%, respectively, while improving average benchmark performance over alternative settings.

Motivation and problem statement

Continued pre-training (CPT) adapts a pre-trained LLM to domain-specific corpora, and its outcome is highly sensitive to the learning rate ($LR$) and batch size ($B$). In practice, these hyperparameters are selected through heuristics or grid searches executed directly on the target model, which is computationally prohibitive at scale. While scaling laws relating loss to compute are well established for training from random initialization (Kaplan et al., 2020), their extension to CPT has remained underexplored for two reasons: it is unclear whether predictable relationships between training state and optimal hyperparameters exist when training does not start from scratch, and distribution shift between pre-training and target data makes it difficult to quantify how much "progress" an initial checkpoint represents on the target domain. This paper addresses both problems with a two-stage framework that predicts optimal $LR$ and $B$ zero-shot for any checkpoint (2606.05610).

Empirical discovery of hyperparameter scaling laws

The authors first establish that optimal CPT hyperparameters follow stable, monotonic trends as compute budget grows: optimal batch size increases monotonically with compute, while optimal learning rate decreases. Using proxy models of 100M and 500M activated parameters (both dense and MoE variants), trained on a 55B-token mixed corpus (SlimPajama 82%, OpenWebMath 9%, GitHub 9%), they perform exhaustive grid searches over $B \in \{64,128,256\}$ and $LR \in \{1,2,5\}\times10^{-4}$.

Optimality is defined compute-theoretically: for a target validation loss $\hat{L}$, the optimal configuration minimizes total FLOPs required to reach that loss. Iso-loss curves in the $(B, LR)$ plane identify this vertex, and fitting across loss levels yields power-law mappings $B_{\text{opt}}(L)$ and $LR_{\text{opt}}(L)$. Composing these with a Kaplan-style loss-compute law $B$0 produces direct compute-to-hyperparameter laws. The regression fits are tight (e.g., slope $B$1 for $B$2 vs. $B$3 on dense models), supporting the paper's central claim that CPT inherits predictable scaling structure rather than requiring ad hoc tuning.

A notable architectural finding emerges from the fitted exponents: MoE models show markedly different sensitivity than dense ones. The learning-rate scaling coefficient for MoE ($B$4) is far less negative than for dense models ($B$5), while the batch-size growth rate for MoE ($B$6) more than doubles that of dense models ($B$7). The authors attribute this to sparse activation: each expert requires sufficient token coverage to stabilize routing, implying that very large MoE models may become bottlenecked by batch-level parallelism rather than raw FLOPs. This is a hypothesis, not a demonstrated mechanism, but it carries practical weight for anyone planning large-scale sparse training.

Equivalent pre-training compute

The second contribution addresses non-zero initialization. The framework posits that any checkpoint's state on a target domain can be characterized by its equivalent pre-training compute $B$8: the compute that would have been required to reach the checkpoint's current validation loss when training from scratch on the target data. Concretely, one evaluates the checkpoint's initial validation loss $B$9 on the CPT corpus and inverts the fitted scaling law:

$LR$0

Total effective compute is then $LR$1, and the predicted terminal loss $LR$2 is fed through the hyperparameter mappings to yield $LR$3 and $LR$4. This formulation treats CPT as a continuation of a single optimization trajectory rather than a disjoint phase. The assumption embedded here—that validation loss on the target domain uniquely determines optimization state, regardless of whether the model arrived there via pre-training or via target-domain training—is strong, and its empirical validity under substantial domain shift is precisely what the ablations test.

For the validation runs, the framework predicts $LR$5 for a Dense-8B checkpoint (initial val loss 2.08, $LR$6) and $LR$7 for an MoE-3B checkpoint (val loss 1.25, $LR$8).

Validation results

Controlled ablations decoupling $LR$9 and $B$0 around the predictions confirm optimality. At fixed batch size, an aggressive $B$1 of $B$2 collapses average benchmark performance to 47.6 (Dense-8B) and 50.8 (MoE-3B) versus 51.3 and 57.9 for the predicted settings, with reasoning benchmarks such as GSM8K degrading most sharply—consistent with catastrophic forgetting of pre-trained weight structure. An overly conservative $B$3 also underperforms. Batch size exhibits an inverted U-shape: too small yields high gradient variance; too large reduces update count within a fixed budget. The predicted configurations achieve the best average scores among all tested settings on both architectures.

Two ablation variants isolate the value of $B$4. Ignoring the pre-trained state entirely (treating CPT as from-scratch training) drops average scores by 0.9 points (Dense-8B) and 1.7 points (MoE-3B) and induces initial loss divergence. Substituting raw pre-training token counts for $B$5 inflates the effective compute and predicts an overly conservative configuration, losing 0.1–0.2 points relative to the full method. Neither the CPT budget alone nor raw pre-training scale suffices; the loss-based calibration is what bridges the distribution shift.

On efficiency, replacing grid search on the target model with proxy-model search reduces hyperparameter search cost by 92.0% for Dense-8B ($B$6 → $B$7 FLOPs) and 73.1% for MoE-3B ($B$8 → $B$9 FLOPs). Because proxy cost is largely independent of target scale, the savings grow with model size—an implication that makes the approach increasingly attractive for larger targets.

Limitations and open questions

The authors concede two boundaries. First, validation covers only up to 8B dense and 3B MoE parameters; whether the fitted power-law constants transfer invariantly to 70B+ models, where training instabilities may alter scaling behavior, remains unverified. Second, the scaling laws are fit on a single representative data mixture; highly specialized domains may exhibit different exponents $B \in \{64,128,256\}$0 and irreducible losses $B \in \{64,128,256\}$1, potentially requiring localized recalibration. A further structural caveat is that the equivalent-compute hypothesis rests on validation loss being a sufficient statistic for optimization state—a claim validated empirically here but not proven, and one whose robustness under extreme distribution shift (e.g., adapting a general model to a narrow scientific corpus) is not directly tested.

Conclusion

This work establishes that optimal CPT hyperparameters follow predictable, architecture-dependent scaling laws, and introduces equivalent pre-training compute as a loss-based coordinate for situating arbitrary checkpoints on those laws. Zero-shot prediction matches or exceeds grid-search-selected configurations on Dense-8B and MoE-3B models while cutting search cost by up to 92%. The main open questions concern extrapolation to much larger scales and to narrowly distributed target domains, where the fitted constants may require re-estimation.