AdamW是一种用于机器学习和深度学习的优化算法,用于训练神经网络。它是Adam优化器的一种变体,将权重衰减与自适应学习率更新解耦。该算法由Ilya Loshchilov和Frank Hutter于2017年提出,并已成为训练大型模型(包括Transformer和大型语言模型)的标准选择。通过将权重衰减与基于梯度的参数更新分离,AdamW解决了原始Adam优化器中的一个已知问题,即权重衰减的应用方式会干扰自适应学习率,导致收敛效果不佳和泛化性能下降。
AdamW的主要动机源于观察到在Adam中,L2正则化项(通常用作权重衰减)会被除以梯度平方的指数移动平均的平方根。这种耦合导致有效权重衰减在不同参数和时间上发生变化,从而可能阻碍优化过程。AdamW提出了一种简单的修复方法:在梯度更新之后直接对参数应用权重衰减,使其独立于自适应学习率的缩放。这种解耦已被证明能改善多种任务上的训练性能,包括图像分类和语言建模。
背景:机器学习中的优化
训练神经网络涉及最小化损失函数,通常使用随机梯度下降(SGD)的变体。在SGD中,模型参数通过沿损失函数负梯度方向迭代更新,而梯度是在训练数据的随机子集上计算的。步长(即学习率)控制每次更新的幅度。多年来,研究者提出了许多改进方法以加速收敛并提升最终性能,例如动量、自适应学习率和权重衰减。
权重衰减是一种正则化技术,通过在损失函数中添加与参数平方和成正比的项来惩罚过大的参数值。在标准SGD中,权重衰减等价于L2正则化,但在Adam等自适应方法中,这种等价关系被打破。AdamW的设计初衷正是为了在自适应优化器中恢复权重衰减的预期行为。
Adam优化器
Adam(自适应矩估计)由Diederik Kingma和Jimmy Ba于2014年提出。它通过维护梯度的一阶矩(均值)和二阶矩(未中心化的方差)的指数移动平均,为每个参数维持独立的学习率。更新规则为:
\[ \theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \]
其中\(\hat{m}_t\)和\(\hat{v}_t\)分别是偏差校正后的一阶和二阶矩估计,\(\eta\)是学习率,\(\epsilon\)是为了数值稳定性而加入的小常数。Adam因其对超参数的鲁棒性和在深度网络训练中的快速收敛而广受欢迎。
然而,在原始Adam实现中,权重衰减是通过L2正则化实现的,即在梯度中添加\(\frac{\lambda}{2} \|\theta\|^2\)项。由于Adam会按二阶矩归一化梯度,这导致有效权重衰减变为\(\lambda / \sqrt{\hat{v}_t}\)。这意味着梯度较大的参数会获得较少的正则化,可能导致过拟合和泛化性能下降。
AdamW算法
AdamW通过修改更新规则来解决上述问题:它从梯度中移除L2正则化项,而是在自适应更新之后直接对参数应用权重衰减。更新规则变为:
\[ \theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} - \eta \lambda \theta_t \]
其中\(\lambda\)是权重衰减系数。这种解耦确保了权重衰减对所有参数均匀施加,不受梯度大小的影响。作者认为,这种方法能带来更好的泛化性能和更稳定的训练过程,尤其是在使用较大学习率时。
在论文中,Loshchilov和Hutter证明了AdamW在多个基准任务上优于使用L2正则化的Adam,包括CIFAR-10上的图像分类和Penn Treebank上的语言建模。他们还表明,AdamW对学习率和权重衰减超参数的选择更为鲁棒。
对深度学习的影响
AdamW对人工智能领域产生了深远影响。它现已成为许多主流深度学习框架(如PyTorch和TensorFlow)中的默认优化器,并被广泛用于训练Transformer和大型语言模型。例如,OpenAI、Anthropic和Google DeepMind等机构开发的许多模型都在训练流程中使用了AdamW。
解耦权重衰减在训练大规模模型时尤其有益,因为正则化对于防止在庞大数据集上过拟合至关重要。研究表明,与使用L2正则化的Adam或其他优化器相比,AdamW能提升收敛速度和最终性能。
与其他优化器的比较
AdamW常与其他自适应优化器(如AdaGrad、RMSProp和Adam)进行比较。AdaGrad和RMSProp会根据历史梯度调整学习率,而Adam则结合了动量和自适应学习率。AdamW通过修复权重衰减问题改进了Adam,使其在广泛任务中更加可靠。
另一个相关的优化器是带动量的SGD,它更简单,但通常需要仔细调整学习率调度。AdamW在易用性和性能之间取得了良好平衡,因此成为许多从业者的首选。不过,一些研究表明,如果超参数调整得当,带动量的SGD在特定任务上可能取得更好的泛化性能,但AdamW仍具竞争力,且对超参数选择更为鲁棒。
实际使用注意事项
使用AdamW时,有几个实际注意事项。权重衰减系数\(\lambda\)通常设置为较小值,如0.01或0.1,但可能需要针对具体任务进行调整。学习率通常设置为1e-4或3e-4,尤其是在训练Transformer时。此外,AdamW通常受益于学习率预热调度,即在前几千步内将学习率从较小值逐渐增加到目标值。
在实践中,AdamW已在大多数深度学习库中实现,用户只需指定优化器并设置权重衰减参数即可。例如,在PyTorch中,可以使用torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)。这种简便性促进了其广泛采用。
扩展与变体
研究者提出了多种AdamW的扩展和变体。例如,将AdamW与Lookahead技术结合,后者维护一组较慢更新的参数以实现更平滑的优化轨迹。另一种变体是使用余弦退火学习率调度的AdamW,已在图像分类任务中显示出性能提升。
在大型语言模型的背景下,AdamW常与混合精度训练和梯度累积结合使用,以处理大批量训练。一些框架还提供了AdamW的融合实现,以减少内存占用并提高计算效率,这对于训练数十亿参数的模型至关重要。
结论
AdamW已成为机器学习工具包中的基础优化器。通过将权重衰减与自适应梯度更新解耦,它解决了原始Adam中的一个重要缺陷,从而带来更好的泛化性能和更稳定的训练。其简单性和有效性使其成为从图像分类到自然语言处理等众多深度学习应用中的默认选择。随着该领域的持续发展,AdamW仍然是一种可靠且广泛使用的优化方法。
参考文献
- Loshchilov, I., & Hutter, F. (2017). Decoupled Weight Decay Regularization. arXiv:1711.05101.
- Kingma, D. P., & Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXiv:1412.6980.