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 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-1B) demonstrate that LLR achieves up to 1.5x training speedup and outperforms baselines, notably raising average zero-shot accuracy from 47.09% to 49.02%. 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.

• The paper introduces a heavy-tail guided layerwise learning rate (LLR) method based on empirical spectral density analysis that assigns adaptive learning rates per transformer layer.
• It achieves improved convergence, reduced perplexity, and enhanced zero-shot accuracy with a 1.5× training speedup compared to uniform learning rates.
• The study demonstrates minimal tuning overhead and robust performance across various transformer models, validating HT-SR theory for practical LLM optimization.

Heavy-Tail Guided Layerwise Learning Rates for LLMs: A Technical Review

Problem Statement and Motivation

Configuration of the learning rate (LR) fundamentally impacts the convergence rate and generalization ability of LLMs. The prevalent paradigm—using a uniform LR across all layers of a transformer model—ignores pronounced architectural heterogeneity manifested in distinct Hessian spectra, heavy-tailedness, and functional roles across layers (e.g., attention vs. FFN vs. embeddings). Empirical evidence increasingly indicates that this one-size-fits-all approach is suboptimal for LLM pre-training.

Layerwise learning-rate allocation has been explored, e.g., via weight-to-gradient norm ratios (LARS, LAMB), sharpness-based blockwise schedules, or temperature balancing methods. However, existing approaches are either sensitive to the tuning of hyperparameters, fail to outperform a well-tuned uniform LR, or do not directly leverage the spectral heterogeneity intrinsic to transformers.

Methodology: Layerwise Learning Rate (LLR) from HT-SR Theory

This work introduces Layerwise Learning Rate (LLR): a spectral-statistics-driven LR assignment algorithm founded on Heavy-Tailed Self Regularization (HT-SR) [martin2019traditional, mahoney2019traditional] and empirical spectral density (ESD) analysis. At regular intervals during training, the weight matrices of each layer are subjected to power-law (PL) fitting. The resulting PL exponents ($\alpha$) quantify the layer's heavy-tailedness; a smaller $\alpha$ suggests a heavier tail and, as argued by HT-SR theory, better optimization and generalization characteristics.

LLR exploits the observation that transformer blocks exhibit stable yet heterogeneous $\alpha$ values across the network (Figure 1):

Figure 2: Comparison of ESD distributions across layers of LLaMa-135M under different training methods; LLR balances heavy-tailed properties and improves perplexity from 17.86 (Uniform) to 17.03.

The core LLR procedure is as follows:

• Periodically compute the ESD for each layer's weight correlation matrix, fit a PL, and extract $\alpha$.
• Use a bounded linear scaling function to assign per-layer LRs: higher $\alpha$ yields higher LR (faster adaptation for poorly trained, weakly heavy-tailed layers), lower $\alpha$ yields lower LR (preventing over-updating well-trained, strongly heavy-tailed layers).
• Specific architectural adaptations are introduced: (1) embedding/output layers with persistently high $\alpha$ are assigned the upper bound LR, and (2) a "soft switch" mechanism linearly interpolates LRs to avoid abrupt transitions.

Figure 4: Left: Embedding gains of LLR vs. baselines across LRs; middle: Hard vs. soft switch comparison; right: LLR phase duration vs. perplexity and compute cost.

This pipeline dynamically balances spectral properties during pre-training, as shown by the decrease of intra-layer $\alpha$ standard deviation and convergence to a globally stable heavy-tailed regime.

Experimental Results

The empirical evaluation uses diverse transformer models (LLaMA series from 60M to 3B parameters, GPT-nano), multiple optimizers (AdamW, Muon), and both pre-training and fine-tuning tasks. Extensive ablations and robustness studies are included.

Pre-training Performance and Downstream Generalization

LLR consistently reduces validation perplexity beyond well-tuned strong baselines. For example, on the FineWeb dataset, the perplexity of LLaMa-135M is reduced from 17.86 (Uniform) to 17.03 (LLR), and LLaMA-1B achieves 9.59 vs. 9.77 (Uniform). Larger-scale LLaMA-3B at 30B tokens achieves comparable gains. LLR also leads in zero-shot accuracy: e.g., LLaMA-1B improves average accuracy from 47.09% (Uniform) to 49.02% (LLR) across seven commonsense reasoning benchmarks.

Figure 5: Learning rate sensitivity for different layerwise LR methods: only LLR robustly outperforms Uniform across all LR tunings and models.

Figure 3: Training loss for LLaMa-135M and LLaMa-350M, Uniform vs. LLR, AdamW: LLR attains similar loss with 1.5× fewer tokens.

LLR demonstrates a consistent 1.5× training speedup over Uniform (Figure 3), matching the downstream performance of Uniform with significantly reduced compute budget.

Ablation Studies and Analysis

• Spectral Analysis: LLR systematically equalizes the distribution of $\alpha$ values, balancing learning load across attention and FFN layers, thus reducing over/under-training heterogeneity (Figure 1).
• Robustness to Assignment Schemes: Linear scaling achieves the best outcomes, outperforming log- and sqrt-based scaling functions.
• Efficiency: Computational overhead is minimized by restricting full spectral analysis to the early training phase, with negligible impact on final model performance.
• Architectural Generality: LLR improves perplexity even for homogeneous attention-only architectures and smaller vision transformer backbones (ViT-Tiny).

Figure 8: Training dynamics of LLaMA-350M with AdamW, showing improvements in both loss and zero-shot accuracy.

Comparison to Related Layerwise/Tuning Methods

LLR outperforms LARS, LAMB, sharpness-driven assignments, Adammini, MuP-AdamW, CompleteP, and other HT-SR-based controls (alphadecay, tempbalance) on both pre-training and fine-tuning benchmarks. Particularly notable is the robustness: other layerwise assignments require aggressive hyperparameter search or are only competitive versus untuned Uniform setups, while LLR can inherit the global LR from the Uniform baseline with minimal tuning.

Theoretical and Practical Implications

The results confirm several key hypotheses:

• Layerwise heterogeneity in spectral statistics (quantified by PL exponent $\alpha$) is persistent and significant in LLMs.
• Adaptive per-layer LR assignment grounded in HT-SR theory can directly leverage this heterogeneity for practical optimization gains without complex tuning or prior architectural knowledge.
• Substantial generalization and convergence gains can be achieved by targeting the heavy-tailed spectral distribution directly, rather than proxies such as weight or gradient norms.

Practically, LLR introduces negligible tuning overhead, is compatible with standard LR decay schedules and modern optimizers, and can be seamlessly integrated into existing pre-training pipelines. The method's minimal computational overhead, especially when restricting full ESD computations to the early training phase, makes it feasible for large-scale deployments.

From a theoretical perspective, this work further validates the predictive and prescriptive power of HT-SR theory for deep neural optimizability, generalization, and scaling.

Future Directions

Potential extensions and open avenues include:

• Scaling to even larger models and multi-modal transformers: Future work should evaluate LLR on 10B+ parameter models and across more complex, multi-component architectures.
• Combining with advanced optimizer dynamics: Investigating synergistic effects when pairing LLR with other adaptive scheduling or mixed-precision optimization strategies.
• Online spectral diagnostics: Investigating adaptive intervals or model-specific heuristics for ESD analysis to further optimize compute-accuracy trade-offs.
• Fine-grained per-component adaptation: Beyond layerwise LRs, controlling other hyperparameters (e.g., weight decay, momentum) using spectral metrics.
• Broader application domains: Exploring the application of heavy-tail guided schedule design in domains beyond language, e.g., vision transformers and graph transformers.

Conclusion

This work demonstrates that a uniform learning rate across transformer layers is fundamentally suboptimal for LLM pre-training. Leveraging empirical heavy-tailedness as a spectral diagnostic, the LLR method assigns adaptive, theoretically principled layerwise learning rates, resulting in consistently improved convergence, generalization, and downstream reasoning performance over established baselines. The low-tuning-overhead, practical deployability, and strong empirical results position LLR as an immediately applicable advance for the optimization of large-scale transformer models (2605.22297).