深度学习中的优化器之SGD with Layer-wise Adaptive Rate (LAR) 算法原理与实现细节
题目描述
本文将详细介绍深度学习优化器SGD with Layer-wise Adaptive Rate (LAR) 的原理和实现。这是一个在训练非常深或大规模神经网络时，用于提升训练稳定性和收敛速度的优化技术。它的核心思想是为网络中的不同层自适应地调整学习率，而不是使用全局统一的学习率。这种方法有助于解决在深度神经网络中，由于不同层的梯度量级或参数更新行为差异巨大而导致的优化困难。
循序渐进讲解
第一步：全局统一学习率的问题
在标准的随机梯度下降（SGD）及其变种（如带动量的SGD）中，我们通常为所有模型参数设置一个全局学习率。在深度神经网络中，这可能会导致问题：
- 梯度量级差异大：网络不同层的权重，其梯度的幅度（范数）可能相差几个数量级。例如，靠近输出的层梯度可能很大，而靠近输入的层梯度可能很小。如果使用相同的学习率，可能导致某些层更新过快而不稳定，另一些层更新过慢而收敛停滞。
- 训练不稳定：在大批量训练或使用大学习率时，梯度量级大的层容易导致参数更新步伐过大，引发损失值剧烈震荡甚至溢出（梯度爆炸），使得训练难以稳定收敛。
因此，我们需要一种机制，能够根据每个参数（或每组参数，如每一层）自身的特性来调整其更新步伐。
第二步：核心思想——层间自适应
LAR算法提出一个直观的解决方案：
对于模型中的每一层，其参数更新的步长（即学习率与该层梯度向量的乘积）应该与该层参数的范数成一定比例。
换句话说，对于参数更大的层（通常对应更复杂的变换），我们允许它以更大的幅度更新；对于参数较小的层，我们限制其更新幅度，使其更精细地调整。具体实现中，这种“比例”是通过一个“信任系数”来控制的，但更常见的LAR实现是直接标准化每一层的更新步长。
一个更稳定、更广泛使用的LAR形式（如论文《Large Batch Training of Convolutional Networks》中所用）的核心是：
在计算参数更新向量（梯度）时，先对其进行归一化（缩放为单位范数），然后再乘以一个全局学习率和该层参数的一个缩放系数。
第三步：算法流程详解
假设神经网络有 \(L\) 层，第 \(l\) 层的参数（权重）记为 \(W^{(l)}\)，其梯度记为 \(G^{(l)} = \nabla_{W^{(l)}} \mathcal{L}\)，其中 \(\mathcal{L}\) 是损失函数。全局学习率记为 \(\eta\)。
LAR 的参数更新步骤如下：
-
计算层参数的范数与梯度范数之比：
对于第 \(l\) 层，计算其权重参数的范数 \(\| W^{(l)} \|\) 和梯度范数 \(\| G^{(l)} \|\)。
LAR的一个关键量是 相对梯度范数（Relative Gradient Norm），但它常常直接用“层权重范数”作为一个参考基准。更常用的实际更新规则是对梯度进行裁剪/缩放，使其更新步长与权重范数成比例。 -
计算自适应学习率系数：
定义一个局部学习率系数 \(\lambda^{(l)}\) 给第 \(l\) 层。常见的定义是：
\[ \lambda^{(l)} = \eta \times \gamma \times \frac{\| W^{(l)} \|}{\| G^{(l)} \| + \beta \times \| W^{(l)} \| + \epsilon} \]
其中：
* $ \eta $ 是全局学习率。
* $ \gamma $ 是一个可调的**信任系数**（trust coefficient），通常是一个很小的值（如 0.001），它决定了我们多大程度上相信梯度方向。$ \gamma $ 越小，更新步长相对于参数范数越小，训练越保守。
* $ \| W^{(l)} \| $ 是第 $ l $ 层权重的 L2 范数。
* $ \| G^{(l)} \| $ 是第 $ l $ 层梯度的 L2 范数。
* $ \beta $ 是一个小的常数（如 0.000001），用于防止当 $ \| W^{(l)} \| $ 很小时分母过小。
* $ \epsilon $ 是一个极小值（如 1e-8），用于保证数值稳定性。
**这个公式的含义是什么？**
当 $ \| G^{(l)} \| $ 远大于 $ \beta \times \| W^{(l)} \| $ 时（即梯度很大），$ \lambda^{(l)} \approx \eta \times \gamma \times \frac{\| W^{(l)} \|}{\| G^{(l)} \|} $。这意味着**更新步长**（$ \lambda^{(l)} \times \| G^{(l)} \| $）近似为 $ \eta \times \gamma \times \| W^{(l)} \| $。也就是说，**该层的参数更新量（步长）被限制在与其自身参数范数成正比的一个值上**，这有效地防止了任何一层的更新幅度失控。
- 执行参数更新：
最终，第 \(l\) 层的参数更新公式为：
\[ W_{\text{new}}^{(l)} = W^{(l)} - \lambda^{(l)} \times G^{(l)} \]
或者，更简洁地写成：
\[ W_{\text{new}}^{(l)} = W^{(l)} - \eta \times \gamma \times \frac{\| W^{(l)} \|}{ \| G^{(l)} \| + \beta \times \| W^{(l)} \| + \epsilon} \times G^{(l)} \]
可以看到，**更新方向**仍然是梯度 $ G^{(l)} $ 的负方向，但**更新步长**不再是简单的 $ \eta $，而是被自适应系数 $ \lambda^{(l)} $ 调整了。
第四步：直观理解与效果
- 对梯度量级大的层：梯度范数 \(\| G^{(l)} \|\) 大，导致自适应系数 \(\lambda^{(l)}\) 变小。这相当于自动降低了该层的学习率，抑制了其更新的幅度，防止不稳定。
- 对梯度量级小的层：梯度范数 \(\| G^{(l)} \|\) 小，导致自适应系数 \(\lambda^{(l)}\) 相对变大。这相当于自动提高了该层的学习率，使其能以更合适的步伐前进，避免收敛过慢。
- 对参数范数大的层：权重范数 \(\| W^{(l)} \|\) 大，分子变大，允许其有相对更大的更新空间，因为参数本身数值大，需要更大的变化才能产生显著影响。
最终效果：LAR 使得网络中所有层的参数更新“步长”（以L2范数衡量）被规范化到与其自身“规模”（权重范数）相适应的水平。这极大地提升了训练的稳定性，特别是在使用非常大批量（Large Batch）进行训练时，使得使用更大的全局学习率 \(\eta\) 成为可能，从而加速收敛。
第五步：实现要点与注意事项
- 与LARS优化器的关系：LAR 的核心思想与 Layer-wise Adaptive Rate Scaling (LARS) 非常相似，甚至经常被混用。但通常，LARS 是 LAR 思想的一种更流行的具体实现，尤其在计算机视觉领域。它们的精神内核一致：逐层自适应调整更新步长。
- 与动量（Momentum）结合：在实际使用中，LAR 通常与动量（Momentum）结合。此时，计算自适应系数 \(\lambda^{(l)}\) 时，范数 \(\| G^{(l)} \|\) 应该替换为动量的范数 \(\| M^{(l)} \|\)，其中 \(M^{(l)}\) 是该层动量缓冲区的值。即先计算带梯度的动量，再对动量应用LAR缩放。
- 参数分组：通常按网络层（如卷积层、全连接层）对参数进行分组，每组应用独立的LAR系数计算。偏置（bias）参数和归一化层（如BatchNorm）的参数通常不应用LAR，因为它们自身特性与权重不同，通常使用全局学习率。
- 超参数调优：关键超参数是信任系数 \(\gamma\)。它控制了更新步长与参数范数之间的比例。通常需要通过实验来找到一个合适的值（如 0.001, 0.005 等）。
总结
SGD with Layer-wise Adaptive Rate (LAR) 通过引入一个与层权重范数和梯度范数之比相关的自适应系数，动态地为神经网络的不同层调整有效的学习率。其核心优势在于稳定了训练过程，特别是当网络很深、参数量级差异大或使用大批量训练时。它使得优化器能够在不导致梯度爆炸的前提下，更激进地更新参数，从而加速模型收敛，是大规模分布式深度学习训练中一项重要的优化技术。