Technical Summary: GQA-µP: The Maximal Parameterization Update for Grouped Query Attention

Problem Statement

Hyperparameter transfer across model architectures, specifically the Maximal Update Parameterization (µP), allows for the determination of optimal hyperparameters for large-scale models by sweeping small proxy models. While µP has been successfully applied to learning rate transfer across model widths, its application to novel architectures and other hyperparameters (such as weight decay and depth) remains challenging.

A specific gap exists for Grouped-Query Attention (GQA), a mechanism widely used in modern Large Language Models (LLMs) to reduce memory usage by sharing key and value heads across groups of query heads. Previous attempts to apply standard µP to GQA revealed a critical failure: while the implementation passed standard "coordinate checks" (validating activation norms), it failed to transfer learning rates effectively across different numbers of KV heads. Furthermore, existing µP frameworks struggled to derive principled scalings for weight decay and depth without relying on "lazy-learning" assumptions. The core theoretical difficulty arises because GQA weight matrices are often rank-degenerate (low-rank), causing the standard spectral norm to fail in accurately capturing the expected deformation of inputs during training.

Methodology

The authors build upon the spectral feature-learning view of Yang et al. (2023a) to derive new scaling laws. The methodology involves three primary theoretical and analytical advancements:

1. Redefining Feature Learning via Spectral Norms:
   The paper elevates spectral norm conditions on weight matrices from a heuristic to the formal definition of feature learning. The authors argue that ensuring the spectral norm of weight updates scales correctly is a stronger and more reliable condition for feature learning than analyzing activation norms alone. This shift allows for the derivation of Complete-P depth and weight-decay scalings without recourse to lazy-learning assumptions.

2. Introduction of the Expected Operator Norm:
   To address the rank-degeneracy of GQA matrices, the authors introduce a modified norm: the expected operator norm ($||A||_E$).

   • Standard spectral norms ($||A||$) measure the maximal deformation of an input vector. For full-rank matrices, this is attained asymptotically.
   • For rank-degenerate matrices (like GQA), the vectors causing maximal deformation occur with probability zero during training (as inputs are i.i.d.). Consequently, the spectral norm overestimates the actual deformation.
   • The expected operator norm is defined as the expectation of the ratio of output to input norms over the input distribution: $||A||_{E,\Omega,p} := \mathbb{E}_{x \sim \Omega} [\frac{||Ax||_p}{||x||_p}]$.
   • This norm correctly reflects the actual deformation encountered during training, restoring the validity of scaling laws for GQA.

3. Derivation of GQA-µP Scalings:
   Using the expected operator norm, the authors derive the specific scaling factors for GQA layers.

   • For a GQA layer with $r$ key-value head repetitions, the learning rate ($\eta$) and weight decay ($\lambda$) must be scaled differently than in standard attention.
   • The derived learning rate scaling is $\eta = \Theta(\frac{1+\sqrt{r}}{n})$, where $n$ is the embedding dimension.
   • The weight decay scaling is adjusted to $\lambda = \Theta(\frac{2n}{1+\sqrt{r}})$ to maintain the balance between regularization and gradient updates.
   • These derivations are summarized in Table 1 of the paper, providing a unified parameterization for Transformers with GQA.

4. Depth Scaling (Complete-P):
   The framework is applied to residual networks to derive depth scalings. By applying the spectral condition to the composition of residual blocks, the authors show that the scaling factor $\beta$ for residual branches must be $\Theta(L^{-1})$ (where $L$ is depth) to prevent signal explosion or vanishing, recovering the same results as Dey et al. (2025) without requiring the "no lazy learning" assumption.

Key Contributions

1. First µP Derivation for GQA: The paper provides, to the authors' knowledge, the first principled derivation of µP scalings for Grouped-Query Attention, resolving the discrepancy where standard µP implementations failed to transfer learning rates.
2. Theoretical Extension of Spectral µP: The work extends the spectral µP theory of Yang et al. (2023a) by promoting spectral norm conditions to the definition of feature learning and introducing the expected operator norm to handle rank-degenerate weights. This enables the derivation of scalings for weight decay, recursion blocks, and GQA.
3. Empirical Validation of Transfer: The authors demonstrate that their derived scalings enable the transfer of learning rates across different GQA repetition counts ($r$) and weight decay across model sizes. They show that the vanilla Adam-µP implementation fails spectral coordinate checks for GQA, whereas their proposed scaling passes them.

Experimental Results

The authors validate their theoretical derivations through extensive experiments:

• Learning Rate Transfer: Experiments varying the number of KV heads (from 1 to 12) show that the standard parameterization and vanilla Adam-µP fail to converge to a consistent optimal learning rate. In contrast, the proposed GQA-µP scaling narrows the optimal learning rate region significantly, enabling effective transfer.
• Coordinate Checks: Analysis of activation update norms ($||\Delta h||$) suggested the vanilla implementation was correct, but analysis of weight matrix spectral norms ($||\Delta W||$) revealed a non-linear dependency on the number of KV heads, explaining the transfer failure. The proposed GQA-µP scaling removes this dependency.
• Weight Decay and $\tau_{\text{epoch}}$ Transfer: The paper demonstrates that with the correct scalings, both weight decay and the training-time constant $\tau_{\text{epoch}}$ (introduced by Wang & Aitchison, 2024) are transferable across model sizes. The variance in optimal hyperparameters is significantly lower for the proposed implementation compared to baselines.
• Noise in GQA: The authors observe that while learning rate transfer is possible, GQA training introduces inherent noise that becomes more pronounced as the number of KV heads decreases. They note a trade-off: achieving feature learning via the derived scaling leads to noisier dynamics as $r$ decreases, while constraining variance to stabilize training shifts the optimal learning rate.

Significance and Claims

The paper claims to close a significant gap in the applicability of µP to modern, production-grade LLM architectures. By identifying the limitations of the standard spectral norm for low-rank matrices and proposing the expected operator norm, the authors provide a robust theoretical foundation for deriving µP scalings for complex architectures like GQA.

The significance of this work lies in:

• Practical Utility: It enables practitioners to determine optimal hyperparameters for large GQA-based models (e.g., LLaMA 3, Mistral 7B) using small proxy models, reducing the computational cost of hyperparameter tuning.
• Theoretical Rigor: It refines the definition of feature learning, moving beyond activation-based checks to weight-matrix-based spectral conditions, which proves more effective for validating implementations.
• Broad Applicability: The framework successfully recovers existing scalings for weight decay and depth (Complete-P) while extending them to GQA, suggesting a unified approach for future architectural innovations.

The authors modestly note that while their asymptotic analysis is sound, the practical range of the repetition parameter $r$ is finite and typically small (e.g., $r \le 16$). The primary goal of the scaling with respect to $r$ is empirical: to prevent learning rate drift when changing the number of repetitions, rather than to model an infinite limit. They also highlight that GQA training remains inherently noisy, suggesting caution when transferring learning rates across vastly different GQA configurations.