Learning Rate Transfer in Normalized Transformers

Abstract: The Normalized Transformer, or nGPT (arXiv:2410.01131) achieves impressive training speedups and does not require weight decay or learning rate warmup. However, despite having hyperparameters that explicitly scale with model size, we observe that nGPT does not exhibit learning rate transfer across model dimension and token horizon. To rectify this, we combine numerical experiments with a principled use of alignment exponents (arXiv:2407.05872) to revisit and modify the $μ$P approach to hyperparameter transfer (arXiv:2011.14522). The result is a novel nGPT parameterization we call $ν$GPT. Through extensive empirical validation, we find $ν$GPT exhibits learning rate transfer across width, depth, and token horizon.

• The paper introduces vGPT, a novel parameterization that ensures robust learning rate transfer in normalized Transformers across width, depth, and token horizon.
• The paper demonstrates that standard pP methods fail for normalized architectures, with vGPT leveraging measured alignment exponents for stable and efficient optimization.
• Extensive ablations confirm that vGPT maintains performance with near-perfect transfer, paving the way for refined hyperparameter scaling protocols.

Learning Rate Transfer in Normalized Transformers: An Authoritative Analysis

Overview

"Learning Rate Transfer in Normalized Transformers" (2604.27077) addresses the principled selection and scaling of learning rates in Normalized Transformer architectures, specifically nGPT, across model width, depth, and token horizon. The paper identifies deficiencies in standard pP-parameterization-based hyperparameter transfer methods for nGPT and proposes a novel parameterization, vGPT, informed by empirical measurement of alignment exponents and theoretical desiderata for feature evolution under optimization. The results demonstrate robust learning rate transfer in vGPT without deterioration in model performance, substantiated by extensive ablations and scaling sweeps.

Normalized Transformers and the Hyperparameter Transfer Problem

Normalized Transformers (nGPT) utilize strict normalization of weights and activations throughout the residual stream, rendering weight decay and learning rate warmup unnecessary. Normalization is applied to both residual and attention projections, combined with trainable scaling vectors ($\mathbf{Q}_A$, $\mathbf{Q}_M$, etc.) that modulate the forward pass via component-wise interpolation. These architectural features reduce the complexity of hyperparameter interaction but introduce new scaling behavior, especially pertaining to width/depth expansion and token horizon.

The core challenge addressed is practical—transferring hyperparameters tuned at small scale to models with much larger depth/width or longer token horizons. Existing pP-type approaches, developed under different alignment assumptions, do not yield reliable transfer for normalized architectures, as empirically observed in nGPT.

The vGPT Parameterization: Derivation and Theoretical Justification

The authors construct vGPT by reconciling theoretical desiderata of stability and feature learning with measured alignment exponents in nGPT. Two main desiderata underpin the analysis:

• Stability at Initialization: Pre-activations and logits maintain order-1 scale at initialization.
• Stable Non-trivial Feature Learning: Updates to preactivations remain order-1 during training, avoiding blowup or triviality.

Alignment exponents quantify the statistical scaling between weight updates and activation changes—full alignment gives order-1 scaling, no alignment yields scaling with $d_{\text{model}}^{-1/2}$. Empirical measurement reveals that weight-activation alignment is consistently mid—i.e., exponent $3/4$—not full or none, across training. This contradicts assumptions in classical pP parameterization, especially regarding $\mathbf{Q}_A$, $\mathbf{Q}_M$ scaling.

Based on these findings, vGPT introduces the following learning rate multipliers:

• Embedding Weights: Scaled by $m_{\text{width}}^{-1/2}$
• Hidden/Output Weights: Scaled by $m_{\text{width}}^{-3/4}$
• Global Learning Rate: Scaled by $m_{\text{data}}^{-1/3}$ in accordance with power-law fits to optimal learning rate versus token horizon from Bjorck et al. (2025) and independent ablations
• Depth Corrections: Initial interpolation parameters ($\mathbf{Q}_A$, $\mathbf{Q}_M$0) scaled linearly with $\mathbf{Q}_M$1

These choices optimize for maximal stable update without breaking feature learning across scaling axes. Notably, depth corrections are minimal, justified by empirical stability and theoretical analysis of the normalization’s effect on update accumulation.

Empirical Validation: Strong Numerical Evidence

The paper's numerical experiments span width sweeps (number of heads), depth sweeps (number of blocks), and token horizon scaling, on large Transformer models with up to 3.2B parameters and over 100M tokens in FineWeb-Edu. Key findings include:

• Learning Rate Transfer Across Width: vGPT achieves near-perfect transfer of optimal learning rate across width, confirmed by stable validation loss curves and minimal drift. In contrast, baseline nGPT and pP-style width corrections exhibit instability and non-transfer.
• Learning Rate Transfer Across Depth: Both vGPT and baseline nGPT show transfer over depth; vGPT displays increased stability at extreme depths and lower sensitivity to learning rate. Depth correction in hidden learning rate (as in Depth-pP, $\mathbf{Q}_M$2) breaks transfer, suggesting depth correction is inappropriate for normalized models with trainable interpolation parameters.
• Token Horizon Scaling: Optimal learning rate decreases proportional to $\mathbf{Q}_M$3, matching previous power-law fits for unnormalized architectures and reaffirming empirical scaling derived in this work.

Validation losses across scaling sweeps are quantitatively tight, with negligible performance drop relative to baseline at any scale, fully matching or exceeding prior best-known results.

Alignment Exponents and Logit Scaling

A substantial contribution is the detailed empirical measurement and theoretical reconciliation of alignment exponents. The paper asserts—with substantiating data—that the mid alignment regime ($\mathbf{Q}_M$4) fits observed Transformer behavior, not full or no alignment. This leads to:

• Logits scaling order-1 at initialization and during training, in contrast to pP which initializes logits at $\mathbf{Q}_M$5 scale.
• Corrections to weight scaling prescriptions, resulting in superior learning rate transfer at scale.

Ablations further show that tuning the learning rate for output weights ($\mathbf{Q}_M$6) is highly sensitive and impactful on performance, but does not affect learning rate transfer, only absolute performance.

Implications and Future Directions

Practically, vGPT enables reliable and performant hyperparameter transfer at arbitrary scale for normalized Transformers, eliminating the need for exhaustive HP search at large scale/fixed compute. This facilitates efficient deployment and exploration in settings with limited computational budget.

Theoretically, the work challenges and refines the prevailing hyperparameter transfer paradigm, emphasizing the relevance of alignment exponent measurement and data-driven correction of classical theories. The authors highlight several unresolved theoretical directions:

• Dynamic analysis of alignment exponents during training, especially at scale
• Analytical derivation of depthwise transfer phenomena in normalized networks
• Effect of weight decay in normalized models and its theoretical reconciliation with learning rate transfer
• Precise analytic justification for the $\mathbf{Q}_M$7 exponent governing token horizon scaling

Conclusion

"Learning Rate Transfer in Normalized Transformers" (2604.27077) establishes vGPT, a new parameterization for normalized Transformers that exhibits robust learning rate transfer across width, depth, and token horizon with empirical and theoretical backing. The paper resolves fundamental inconsistencies in pP-style transfer for normalized architectures, provides strong numerical evidence for its claims, and defines new standards for principled hyperparameter scaling. The methodology and results suggest further refinement of transfer theory and practical scaling protocols is achievable with more precise empirical and theoretical instruments, with implications for efficient training of large-scale LLMs and other deep architectures.