Revealing Modular Gradient Noise Imbalance in LLMs: Calibrating Adam via Signal-to-Noise Ratio

📄 arXiv: 2605.05794v1 📥 PDF

作者: Ziqing Wen, Zhouyang Liu, Jiahuan Wang, Ping Luo, Li Shen, Dongsheng Li, Tao Sun

分类: cs.LG, cs.AI

发布日期: 2026-05-07

💡 一句话要点

提出基于信噪比（SNR）的模块化学习率缩放方法（MoLS），解决大模型训练中的梯度噪声不平衡问题。

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

关键词: 大语言模型 自适应优化器 梯度噪声 信噪比 模型训练优化 深度学习

📋 核心要点

  1. 现有Adam优化器缺乏对LLM异构模块间梯度噪声差异的显式建模，导致训练不稳定或收敛效率低下。
  2. 论文提出MoLS方法，通过实时估计模块级信噪比（SNR）来动态调整学习率，实现自动化模块级优化。
  3. 实验证明MoLS在多个LLM基准上显著提升了收敛速度与泛化性能，且无需昂贵的人工超参数调优。

📝 摘要（中文）

大语言模型（LLM）的卓越性能源于其庞大的规模和异构的模块组成，但这种结构异构性也带来了优化挑战。尽管Adam(W)等自适应优化器提供了参数级的自适应能力，却未能显式处理模块级的梯度异构性，导致收敛缓慢、性能次优或训练不稳定。现有方法通常依赖于人工调优的模块特定学习率，不仅计算成本高昂，且难以在不同任务或模型间泛化。为建立更具原则性的方法，本文分析了Adam在高噪声模块中的噪声抑制行为，并提出了基于信噪比的模块化学习率缩放（MoLS）。MoLS通过估计模块级信噪比来缩放Adam更新，实现了无需人工调优的自动化模块级学习率分配。在多个LLM训练基准上的实验表明，MoLS提升了收敛速度与泛化能力，性能可媲美精心调优的模块特定学习率，且兼容内存高效的训练算法。

🔬 方法详解

问题定义：LLM由多种异构模块（如Attention、MLP、LayerNorm等）组成，各模块的梯度统计特性存在显著差异。Adam优化器虽然具备参数级自适应，但无法感知模块间的梯度噪声不平衡，导致某些模块在训练过程中更新步长不合理，限制了模型性能。

核心思路：论文的核心思想是利用信噪比（SNR）作为衡量梯度质量的指标。通过分析发现，Adam在高噪声模块中存在过度更新的倾向，MoLS通过引入SNR反馈机制，自动抑制高噪声模块的更新幅度，从而实现更平稳的优化过程。

技术框架：MoLS框架嵌入在Adam优化器中，主要包含两个阶段：首先是模块级SNR的实时估计，通过统计梯度的一阶和二阶矩信息计算各模块的信噪比；其次是基于SNR的缩放因子计算，将该因子作用于Adam的更新步长，实现动态学习率调整。

关键创新：MoLS的本质创新在于将“模块级梯度质量”显式引入优化器更新逻辑中。与传统人工设置模块学习率的方法不同，MoLS是完全自动化的，且不增加额外的计算开销，能够自适应地应对不同模型架构的异构性。

关键设计：该方法通过移动平均法平滑估计梯度均值与方差，进而计算SNR。设计中引入了缩放函数，将SNR映射为学习率乘数，确保在噪声较大的模块中自动减小步长，在信号清晰的模块中保持或增大步长，从而实现全局最优的收敛路径。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MoLS在多个主流LLM训练基准上表现优异，在无需人工干预的情况下，收敛速度较标准AdamW有显著提升。实验结果显示，其性能表现与经过精细调优的模块特定学习率策略相当，且在内存高效训练（如LoRA或ZeRO）中表现出良好的兼容性，证明了该方法在复杂模型训练中的鲁棒性与高效性。

🎯 应用场景

该方法适用于大规模语言模型（LLM）的预训练与微调阶段，特别是在模型架构复杂、模块异构性强的场景下。其自动化特性使其在资源受限的训练环境中极具价值，可显著降低超参数搜索成本，并提升模型在下游任务中的泛化能力与训练稳定性。

📄 摘要（原文）

The impressive performance of large language models (LLMs) arises from their massive scale and heterogeneous module composition. However, this structural heterogeneity introduces additional optimization challenges. While adaptive optimizers such as Adam(W) provide per-parameter adaptivity, they do not explicitly account for module-level gradient heterogeneity, resulting in slower convergence, suboptimal performance, or training instability. Existing approaches typically rely on manually tuned module-specific learning rates or specific optimization strategies, which are computationally costly and difficult to generalize across tasks or models. To establish a more principled approach, we first analyze the noise-damping behavior of Adam in high-noise modules and introduce \textbf{Module-wise Learning Rate Scaling via SNR (MoLS)}. MoLS estimates module-level SNRs to scale Adam updates, allowing automated module-wise learning rate allocation without manual tuning. Empirical results through multiple LLM training benchmarks demonstrate that MoLS improves convergence speed and generalization, achieving performance comparable to carefully tuned module-specific learning rates, while remaining compatible with memory-efficient training algorithms.