Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Duration (2512.22382v1)

Abstract: Hyperparameter tuning can dramatically impact training stability and final performance of large-scale models. Recent works on neural network parameterisations, such as $μ$P, have enabled transfer of optimal global hyperparameters across model sizes. These works propose an empirical practice of search for optimal global base hyperparameters at a small model size, and transfer to a large size. We extend these works in two key ways. To handle scaling along most important scaling axes, we propose the Complete$^{(d)}$ Parameterisation that unifies scaling in width and depth -- using an adaptation of CompleteP -- as well as in batch-size and training duration. Secondly, with our parameterisation, we investigate per-module hyperparameter optimisation and transfer. We characterise the empirical challenges of navigating the high-dimensional hyperparameter landscape, and propose practical guidelines for tackling this optimisation problem. We demonstrate that, with the right parameterisation, hyperparameter transfer holds even in the per-module hyperparameter regime. Our study covers an extensive range of optimisation hyperparameters of modern models: learning rates, AdamW parameters, weight decay, initialisation scales, and residual block multipliers. Our experiments demonstrate significant training speed improvements in LLMs with the transferred per-module hyperparameters.

• The paper introduces a Complete parameterization that robustly transfers optimizer hyperparameters across key scaling axes such as model width, depth, batch size, and token horizon.
• It employs SDE-based scaling rules for AdamW and QK normalization, ensuring stability and accelerated convergence in Transformer architectures.
• Per-module hyperparameter optimization using trust-region search yields practical performance gains, effectively scaling from small to large models.

Completed Hyperparameter Transfer Across Model Structure, Batch, and Token Horizon

Introduction

This paper provides a systematic framework for hyperparameter transfer in large-scale neural network training, generalizing and extending the principles of parameterization such as $\mu$P and its variants. The core advancements are the introduction of the Complete parameterization, which unifies and enables robust transfer of optimizer hyperparameters (learning rate, weight decay, Adam parameters, initialization scale) across key scaling axes: model width, depth, batch size, and token horizon. The authors specifically address the empirical and theoretical challenges associated with transferring not only global hyperparameters but also highly granular per-module settings, with rigorous validation in modern Transformer architectures.

Figure 1: Trust-region-based per-module HP optimization at 50M scale leads to a $1.32\times$ speed-up after transfer to large scale, surpassing global HP baselines.

Theoretical Advances in Parameterization

The Complete parameterization builds upon Tensor Programs theory, $\mu$P, and CompleteP [dey2025completep], extending to critical architectural components, notably Query-Key normalization, and correcting issues in prior scaling formulas (e.g., AdamW’s $\epsilon$ for embeddings). This yields a parameterization that ensures transferability of hyperparameters by maintaining SDE invariance and consistent limiting dynamics across scaling axes. Notably:

• The scheme corrects and extends depth-dependent scaling ($\alpha \in [0.5,1]$) for residual connections, achieving transfer for both theoretically motivated regimes.
• QK normalization receives dedicated scaling rules, accounting for parameter sharing across attention heads.
• The output layer multiplier is absorbed into learning rate and initialization, facilitating efficient memory usage for large-scale vocabularies.

  Figure 2: Empirical verification of global learning rate transfer across width and depth under Complete parameterization for multiple $\alpha$.

  

Figure 3: Addition of QK-norms enhances transfer stability across depth scaling variations.

Hyperparameter Transfer Beyond Model Size

Batch Size Scaling

The paper rigorously extends SDE-based transfer rules previously derived for Adam to AdamW, including the first theoretically-justified scaling rule for AdamW’s weight decay:

$$\eta' = \sqrt{\kappa}\eta,\quad \lambda' = \sqrt{\kappa}\lambda$$

where $\kappa$ is the batch size increase factor.

This yields near-invariant training trajectories and optimally preserved minimum-loss solutions as a function of batch size, as shown by direct empirical comparison.

Figure 4: Learning rates scale appropriately with batch size under SDE rule; naive transfer fails to preserve optima.

Figure 5: Corrected (SDE-based) scaling of weight decay is essential for successful batch size transfer.

Token Horizon Scaling

For token horizon (number of training tokens), the optimal learning rate scales roughly as $1/\sqrt{\kappa}$ when longer horizons are used at constant batch size, as predicted by SDE analysis. The authors verify that this rule generalizes beyond learning rate, encompassing weight decay and AdamW’s additional hyperparameters.

Importantly, when the token horizon is scaled via batch size—increasing batch at fixed iteration count—transfer is further improved due to a boosted signal-to-noise ratio; scaling laws confirm better asymptotic performance at large compute budgets.

Figure 6: Learning rate transfer across token horizon validated via SDE analysis; confirms optimal $\eta$ scaling.

Figure 7: Scaling law comparison—CompleteP with SDE token-horizon scaling outperforms alternative scaling policies at large model sizes.

Per-Module Hyperparameter Optimization

The study demonstrates that optimizing hyperparameters specifically for each module (and depth index) yields substantial acceleration and loss improvements over global HP settings, even after transfer to much larger models. The optimization landscape is highly irregular, often with “cliffs” where training diverges, diminishing the effectiveness of random search or standard Bayesian optimization. Nonetheless, the authors find it is locally invex, permitting efficient trust-region or evolutionary search strategies.

Figure 8: The stable region boundary for per-module learning rates exhibits high complexity, challenging naive optimization.

Figure 9: Depth-Kronecker parameterization reduces optimization search space dimensionality while nearly saturating performance gains.

Notably, the full uncoupling of per-layer learning rate multipliers yields negligible improvement over the depth-Kronecker factorization, indicating that most benefits of per-module HP optimization are already captured with structured parameterization.

Transferability and Practical Impact

Empirically, the optimal per-module HPs discovered at small scale (e.g., 50M parameters) transfer effectively to models up to 7B parameters with $>100\times$ compute increase, within the Complete transfer regime. The architecture-wide transfer preserves benchmark performance and accelerates convergence in large-scale pretraining scenarios.

Figure 10: Downstream benchmark: per-module HP transfer yields consistent performance gains over global HP transfer at 7.2B scale.

Figure 11: Interpolation between global and per-module HP multipliers confirms robust transfer and improvement across scales.

Figure 12: Transfer penalty when HPs are optimized at excessively small scales is minor above $\sim$136M model size, guiding practical search strategies.

Limitations and Future Directions

While the approach is broadly applicable, model, dataset, and optimizer-specific dependencies may affect transfer rules’ reliability. The study focuses on autoregressive Transformer training with a cosine LR schedule; further research is needed to confirm generality across architectures, domains, and scheduling strategies. Improvements in trial efficiency for per-module HP search, such as trust-region Bayesian optimization or early-stopping strategies, are promising avenues. The diminishing returns of per-module HP gains at extreme scale remain open for theoretical analysis.

Conclusion

This work formalizes and validates a unified parameterization, Complete, enabling principled hyperparameter transfer across model width, depth, batch size, and training horizon, for both global and per-module settings. The methodology is supported by theoretical analysis and large-scale empirical tests, providing a practical recipe for efficiently tuning and scaling deep neural networks. This framework is poised to be influential for efficient LLM pretraining and scaling, subject to further exploration in diverse model regimes.