Spectral Condition for $μ$P under Width-Depth Scaling
This work addresses a critical problem for researchers and practitioners scaling large models by providing a unified and simpler approach to hyperparameter tuning across model sizes, though it is incremental as it builds on existing μP methods.
The paper tackled the challenge of stable feature learning and hyperparameter transfer in generative foundation models scaled in both width and depth by developing a spectral framework for maximal update parameterization (μP) under joint width-depth scaling, demonstrating that it preserves stable feature learning and enables robust hyperparameter transfer in GPT-2 style language models.
Generative foundation models are increasingly scaled in both width and depth, posing significant challenges for stable feature learning and reliable hyperparameter (HP) transfer across model sizes. While maximal update parameterization ($μ$P) has provided a principled solution to both problems for width scaling, existing extensions to the joint width-depth scaling regime remain fragmented, architecture- and optimizer-specific, and often rely on technically involved theories. In this work, we develop a simple and unified spectral framework for $μ$P under joint width-depth scaling. Considering residual networks of varying block depths, we first introduce a spectral $μ$P condition that precisely characterizes how the norms of weights and their per-step updates should scale with width and depth, unifying previously disparate $μ$P formulations as special cases. Building on this condition, we then derive a general recipe for implementing $μ$P across a broad class of optimizers by mapping the spectral constraints to concrete HP parameterizations. This approach not only recovers existing $μ$P formulations (e.g., for SGD and AdamW) but also naturally extends to a wider range of optimizers. Finally, experiments on GPT-2 style language models demonstrate that the proposed spectral $μ$P condition preserves stable feature learning and enables robust HP transfer under width-depth scaling.