GQA-{\mu}P: The maximal parameterization update for grouped query attention

2026-05-14

AI Summary

This work addresses the lack of theoretical support for hyperparameter transfer in grouped-query attention (GQA) architectures by introducing a spectral feature learning perspective. It formally incorporates weight spectral norm constraints into the definition of feature learning and proposes a corrected spectral norm to handle non-full-rank weight matrices. Based on this framework, the study derives, for the first time, a maximal update (μP) scaling rule tailored to GQA, which correctly scales both depth and weight decay without relying on the lazy training assumption. Experimental results demonstrate that learning rates and weight decay settings can be effectively transferred across GQA configurations, substantially reducing the computational cost of hyperparameter tuning in large language models.

Abstract

Hyperparameter transfer across model architectures dramatically reduces the amount of compute necessary for tuning large language models (LLMs). The maximal update parameterization ({\mu}P) ensures transfer through principled mathematical analysis but can be challenging to derive for new model architectures. Building on the spectral feature-learning view of Yang et al. (2023a), we make two advances. First, we promote spectral norm conditions on the weights from a heuristic to the definition of feature learning, and as a consequence arrive at the Complete-P depth and weight-decay scalings without recourse to lazy-learning. Second, we consider a modified spectral norm that preserves the valid scaling law of network weights when weight matrices are not full rank. This enables (to our knowledge, the first) derivation of {\mu}P scalings for grouped-query attention (GQA). We demonstrate the efficacy of our theoretical derivations by showing learning rate transfer across the GQA repetition hyperparameter as well as experiments regarding transfer over weight decay.

Problem

• Grouped Query Attention
• maximal update parameterization
• hyperparameter transfer
• scaling laws
• spectral norm

Innovation

• maximal update parameterization
• grouped-query attention
• spectral norm
• hyperparameter transfer
• scaling laws