Learning Rate Transfer in Normalized Transformers
| Authors | Boris Shigida et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2604.27077 |
| Download | |
| Categories | cs.LG, cs.AI, stat.ML |
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.
Engineering Breakdown
Plain English
undefined
Core Technical Contribution
The core novelty is a mathematically principled reparameterization of the Normalized Transformer (nGPT) called νGPT that enables hyperparameter transfer across model dimensions. Rather than proposing an entirely new architecture, the authors apply alignment exponents—a framework from recent work on feature learning—to identify and correct the scaling mismatches in nGPT's original parameterization. This modified approach extends the classical μP (maximal update parameterization) framework, which worked for RNNs and standard Transformers but failed for nGPT's specific normalization strategy. The technical contribution is showing that with the correct scaling relationships, the same learning rate and hyperparameters can be reliably used across models that differ in width, depth, and training token budget.
How It Works
The paper starts by analyzing why nGPT fails at learning rate transfer despite having explicit model-size scaling in its hyperparameters. The authors conduct numerical experiments to measure alignment exponents—scalars that quantify how the gradient signal changes across different dimensions during training—and use these measurements to identify where nGPT's parameterization deviates from the principled scaling required for transfer. They then reformulate nGPT's layer normalization, initialization, and gradient flow to match the theoretical predictions of alignment exponents, producing νGPT. The key transformation involves adjusting multiplicative factors in attention, residual connections, and parameter initialization so that the effective learning signal remains constant regardless of model size. Finally, they validate νGPT empirically by training models of different widths (e.g., 1B to 7B parameters), depths, and sequence lengths using identical hyperparameters, confirming that training curves and final loss converge to the same target across configurations.
Production Impact
For production ML systems, this enables massive cost savings in hyperparameter tuning and model scaling experiments. Without learning rate transfer, teams must repeatedly tune hyperparameters (especially learning rate schedules) when scaling from a 1B-parameter proof-of-concept to a 70B-parameter production model; with νGPT, you can train a small model at low cost, then apply those exact hyperparameters to the production-scale model with confidence. This directly reduces the wall-clock time and compute budget needed for the expensive scaling phase—critical when each training run costs tens of thousands of dollars. The approach also simplifies the engineering pipeline: fewer hyperparameter searches mean less infrastructure for distributed tuning, simpler experiment tracking, and faster time-to-training-stability. The trade-off is that teams must adopt the nGPT architecture and its associated constraints (e.g., specific normalization schemes), and they lose the flexibility to fine-tune hyperparameters per model size, though the paper suggests this loss is negligible in practice.
Limitations and When Not to Use This
The paper's validation is primarily empirical and focuses on dense Transformer models on standard language modeling tasks; it does not establish strong theoretical guarantees for why νGPT's transfer should hold universally or under what distribution shifts it might fail. The approach assumes that alignment exponents remain stable across training and across model scales, which may break down when using novel optimizers, extreme batch sizes, or very long-context training (e.g., 1M tokens). The paper does not address other critical hyperparameters like batch size or optimizer choice—learning rate transfer alone is insufficient for full reproducibility, and practitioners still need guidance on how batch size impacts training dynamics. Additionally, the work is limited to the Normalized Transformer architecture; the lessons may not transfer to other emerging architectures (e.g., hybrid attention-RNN models, mixture-of-experts), and integration with production systems that rely on tools, checkpointing, and inference optimization around standard Transformers could introduce compatibility friction.
Research Context
This work builds directly on three recent papers: μP (arXiv:2011.14522), which established that certain parameterizations enable learning rate transfer in standard Transformers; alignment exponents (arXiv:2407.05872), which provide a principled diagnostic tool for measuring feature learning alignment; and nGPT (arXiv:2410.01131), the Normalized Transformer architecture that motivated the problem. The paper contributes to the growing research direction on hyperparameter transfer and scaling laws, which has become critical as model sizes grow and training costs explode. It also indirectly strengthens the case for normalization-based architectural innovations by solving a key practical barrier—without learning rate transfer, nGPT's other benefits (no warmup, no weight decay, faster training) are diminished by the need to re-tune for each scale. Future work is likely to extend these alignment exponent-based diagnostic techniques to other new architectures and to investigate whether the νGPT parameterization generalizes to multimodal models, mixture-of-experts, and other emerging model families.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::