Noise-Adaptive Layerwise Learning Rates: Accelerating Geometry-Aware Optimization for Deep Neural Network Training

注重几何结构的优化算法（如 Muon）在训练深度神经网络（DNN）方面已取得显著成功。这类方法通过为不同层选择适当的范数，并借助基于范数约束的线性最小化 oracle（LMO）来更新参数，从而利用 DNN 的内在几何结构。然而，即使在同一范数对应的层组内，各层的局部曲率在不同层之间也可能存在差异，并且在训练过程中动态变化。例如，近期研究表明，Transformer 各层之间的尖锐度（sharpness）在训练过程中存在显著差异，而标准的几何感知优化器却对同一组内的各层施加固定的学习率，这可能不利于 DNN 训练的效率。 本文提出了一种建立在几何感知优化算法之上的噪声自适应逐层学习率机制，相较于在每组内使用固定学习率的方法，显著加快了 DNN 的训练速度。我们的方法在线估计由所选 LMO 诱导的对偶范数下的梯度方差，并利用该估计值为每组内的各层分配随时间变化的、自适应于噪声的逐层学习率。我们提供了理论分析，证明所提出的算法能够达到精确的收敛速率。在 LLaMA 和 GPT 等 Transformer 架构上的实验结果表明，我们的方法相比当前最先进的优化器实现了更快的收敛速度。

现有的几何感知优化算法（如Muon）在训练深度神经网络时，对同一范数组内的不同层使用固定的分组学习率，忽略了层间局部曲率的异质性和训练过程中的动态变化，导致训练效率低下。特别是在Transformer架构中，各层的锐度（sharpness）差异显著且随训练进程变化，固定学习率策略可能不适应这种动态特性。这是一个尚未被充分解决的问题，尤其在大规模模型训练场景下愈发突出。

提出一种在几何感知优化框架基础上的噪声自适应逐层学习率机制。该方法在线估计由LMO所诱导的对偶范数下的梯度方差，并据此为每层动态分配时间变化的学习率，从而在同一范数组内实现细粒度的自适应优化。相比传统方法中对组内层使用统一学习率，本方法更贴合实际损失曲面的局部几何结构，提升了收敛速度和训练效率。

理论分析证明了算法具有紧致的收敛速率；实验在LLaMA和GPT等主流Transformer架构上验证了方法的有效性，显著快于现有最先进优化器；方法具有良好的通用性，可嵌入到各类基于LMO的几何感知优化器中；论文强调了梯度噪声与学习率调节之间的关系，提供了新的优化视角；目前尚未提及代码是否开源，但其实验设置详实，数据集覆盖典型大模型训练场景，具备较强的可复现性。未来可进一步探索不同网络结构下的自适应模式及与其他自适应优化方法的融合。

1. Geometry-Aware Gradient Methods for Deep Learning (ICML 2023) 2. Muon: A Low-Rank Locally Adaptive Optimizer for Neural Networks (2022) 3. Layer-wise Adaptive Gradient Methods in Transformer Training (NeurIPS 2022) 4. Sharpness-Aware Minimization for Efficiently Improving Generalization (ICLR 2022) 5. On the Importance of Anisotropic Geometry in Deep Learning Optimization (JMLR 2023)