One LR Doesn’t Fit All: Heavy-Tail Guided Layerwise Learning Rates for LLMs
Abstract
Learning rate configuration is a fundamental aspect of modern deep learning. The prevailing practice of applying a uniform learning rate across all layers overlooks the structural heterogeneity of Transformers, potentially limiting their effectiveness as the backbone of Large Language Models (LLMs). In this paper, we introduce Layerwise Learning Rate (LLR), an adaptive scheme that assigns distinct learning rates to individual Transformer layers. Our method is grounded in Heavy-Tailed Self-Regularization (HT-SR) theory, which characterizes the empirical spectral density (ESD) of weight correlation matrices to quantify heavy-tailedness. Layers with weaker heavy-tailedness are assigned larger learning rates to accelerate their training, while layers with stronger heavy-tailedness receive smaller learning rates. By tailoring learning rates in this manner, LLR promotes balanced training across layers, leading to faster convergence and improved generalization. Extensive experiments across architectures (from LLaMa to GPT-nano), optimizers (AdamW and Muon), and parameter scales (60M–3B, up to 100B tokens) demonstrate that LLR achieves up to 1.5× training speedup and outperforms baselines, notably raising average zero-shot accuracy from 47.09% to 49.02% for 1B models and from 48.58% to 50.61% for 3B models. A key advantage of LLR is its low tuning overhead: it transfers nearly optimal LR settings directly from the uniform baseline. Code is available at https://github.com/hed-ucas/Layer-wise-Learning-Rate.
Lay Summary
Training large AI models typically forces all internal layers to learn at the exact same speed. However, because these layers perform different functions and learn differently, a "one-size-fits-all" pace is inefficient. We introduce Layerwise Learning Rate (LLR), a method that automatically assigns a customized learning speed to each individual layer. LLR monitors each layer's progress during training: layers that are lagging behind are given a faster pace to catch up, while those already performing well are adjusted more carefully. By balancing the learning process across the entire model, LLR speeds up AI training by up to 1.5x and produces noticeably more accurate models, all without requiring engineers to manually guess the right settings.