🤖 AI Summary
This work addresses the challenge of achieving stable hyperparameter transfer across model scales in complex gated linear architectures such as Gated Delta Networks, where existing parametrizations fail. It presents the first extension of Maximal Update Parametrization (μP) to gated linear networks with structured state transitions. By analyzing the coordinate-wise scaling behavior in forward propagation, gating mechanisms, and recurrent state dynamics, the authors derive theoretically grounded scaling rules tailored to this architecture. The resulting parametrization enables zero-shot learning rate transfer across varying model widths. Empirical evaluations in language model pretraining demonstrate that this approach significantly outperforms standard parametrizations, which exhibit unstable cross-scale transfer performance.
📝 Abstract
Training and scaling Large Language Models demand enormous computational resources, motivating both efficient sub-quadratic architectures and principled hyperparameter tuning methods. While the Maximal Update Parametrization ($μ$P) has enabled zero-shot hyperparameter transfer for standard Transformers, its extension to linear models, particularly those with structured state transitions and complicated architectures, remains largely unexplored. By rigorously propagating coordinate-size estimates through the forward pass, gating mechanisms, and recurrent state dynamics, we derive the scaling rules for Gated Delta Network. Experiments on language-model pre-training confirm that our configurations enable stable learning-rate transfer across model widths under both AdamW and SGD, whereas standard parametrization fails to transfer, validating the correctness and practical utility of our analysis.