One LR Doesn't Fit All: Heavy-Tail Guided Layerwise Learning Rates for LLMs

📄 arXiv: 2605.22297v1 📥 PDF

作者: Di He, Songjun Tu, Keyu Wang, Lu Yin, Shiwei Liu

分类: cs.LG, cs.AI

发布日期: 2026-05-21

🔗 代码/项目: GITHUB


💡 一句话要点

提出重尾分布引导的层级学习率调整方法,提升LLM训练效率与泛化能力

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 层级学习率 重尾分布 自适应学习率 Transformer 优化算法 经验谱密度

📋 核心要点

  1. Transformer结构异构性在LLM中未被充分考虑,现有方法通常采用统一学习率,限制了模型性能。
  2. 论文提出基于重尾自正则化理论的层级学习率(LLR)方法,为不同层自适应分配学习率。
  3. 实验证明,LLR在多种模型和优化器上实现了训练加速和零样本准确率的提升,且调优成本低。

📝 摘要(中文)

本文提出了一种名为层级学习率(LLR)的自适应方案,旨在为Transformer架构中的每一层分配不同的学习率。该方法基于重尾自正则化(HT-SR)理论,通过量化权重相关矩阵的经验谱密度(ESD)来表征重尾性。重尾性较弱的层被赋予更大的学习率以加速训练,而重尾性较强的层则获得较小的学习率。通过这种方式调整学习率,LLR促进了各层之间的平衡训练,从而实现更快的收敛和更好的泛化。在LLaMA到GPT-nano等多种架构、AdamW和Muon等优化器以及60M-1B参数规模的广泛实验表明,LLR实现了高达1.5倍的训练加速,并优于基线方法,显著地将平均零样本准确率从47.09%提高到49.02%。LLR的一个关键优势是其低调优开销:它可以直接从统一基线传递近乎最优的LR设置。

🔬 方法详解

问题定义:现有的大语言模型(LLM)训练过程中,通常对Transformer的每一层都采用相同的学习率。然而,Transformer的各层结构和功能存在差异,统一的学习率无法针对性地优化每一层,导致训练效率低下,模型性能受限。因此,需要一种能够根据不同层特性自适应调整学习率的方法。

核心思路:论文的核心思路是基于重尾自正则化(HT-SR)理论,通过分析权重相关矩阵的经验谱密度(ESD)来量化每一层的重尾性。重尾性反映了层内参数的关联程度和训练难度。重尾性较弱的层,参数关联性低,训练难度小,可以采用较大的学习率加速训练;反之,重尾性较强的层,参数关联性高,训练难度大,应该采用较小的学习率,避免训练不稳定。

技术框架:LLR方法的整体框架包括以下几个步骤:1) 使用预训练或随机初始化的模型;2) 计算每一层权重相关矩阵的经验谱密度(ESD);3) 基于ESD计算每一层的重尾性指标;4) 根据重尾性指标,为每一层分配不同的学习率;5) 使用自适应学习率进行模型训练。

关键创新:LLR方法最重要的创新在于将重尾自正则化理论应用于层级学习率的自适应调整。它通过量化每一层的重尾性,实现了对不同层训练难度的评估,并根据评估结果动态调整学习率。与传统的统一学习率方法相比,LLR能够更好地平衡各层的训练进度,提高训练效率和模型性能。

关键设计:LLR的关键设计包括:1) 使用权重相关矩阵的经验谱密度(ESD)来量化重尾性;2) 设计了基于重尾性指标的学习率分配策略,通常重尾性越强,学习率越小;3) 学习率的调整幅度可以通过一个超参数进行控制,以平衡训练速度和稳定性。论文中使用的具体公式和参数设置可以在论文原文中找到。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,LLR方法在多种模型(LLaMA, GPT-nano等)和优化器(AdamW, Muon)上均取得了显著的性能提升。在参数规模为60M-1B的模型上,LLR实现了高达1.5倍的训练加速,并将平均零样本准确率从47.09%提高到49.02%。更重要的是,LLR具有较低的调优开销,可以直接从统一学习率的基线设置迁移近乎最优的学习率。

🎯 应用场景

该研究成果可广泛应用于各种基于Transformer的大语言模型训练,尤其是在资源受限或需要快速迭代的场景下。通过自适应调整层级学习率,可以显著提升训练效率,降低训练成本,并提高模型的泛化能力。此外,该方法还可以应用于其他深度学习模型,例如计算机视觉领域的Transformer模型。

📄 摘要(原文)

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-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.