随机权重平均(SWA)是一种在深度学习中通过平均神经网络在训练轨迹上多个点的权重来改善泛化能力的模型平均技术。与传统训练仅保留最终权重不同,SWA在训练后期按固定间隔收集权重并计算它们的算术平均值。这一简单过程通常会在损失景观中产生更平坦的解,这些解与更好的测试性能和鲁棒性相关。SWA由Pavel Izmailov、Dmitrii Podoprikhin、Timur Garipov、Dmitry Vetrov和Andrew Gordon Wilson在2018年的论文《Averaging Weights Leads to Wider Optima and Better Generalization》中提出。该方法无需改变底层模型架构或损失函数,因此易于集成到现有训练流程中。
SWA的核心思想是,在使用随机梯度下降(SGD)或其变体训练期间,模型参数会在低损失区域附近振荡。对这些参数进行平均可以平滑振荡,并产生一个位于宽而平坦的最小值处的解。这种平坦最小值被假设比尖锐最小值具有更好的泛化能力,因为它们对扰动以及训练和测试分布之间的偏移不那么敏感。SWA已被证明能在各种架构和任务(包括图像分类、目标检测和语言建模)中提高准确性和校准性能。
机制与算法
SWA分两个阶段运行。在第一阶段,模型使用标准学习率调度(例如余弦退火或阶梯衰减)训练,以接近最小值区域。在第二阶段,训练继续使用恒定或循环学习率,并在每个epoch结束或固定迭代次数后收集权重。SWA权重作为运行平均值更新:n_swa = n_swa + 1;w_swa = (w_swa * (n_swa - 1) + w) / n_swa,其中w是当前权重,w_swa是平均权重。最终模型使用w_swa。
一个关键细节是,平均后必须重新计算批归一化统计量,因为批归一化层的运行均值和方差不会随权重一起平均。在实践中,计算SWA权重后,需要对训练数据执行一次前向传播以更新批归一化统计量。这一步对于使用批归一化的模型(如ResNet)至关重要。
理论动机
SWA的成功通常通过损失景观几何的视角来解释。Izmailov等人的研究及后续工作表明,SGD倾向于收敛到低损失区域边界附近的点,而对多个点进行平均会将解移向该区域的中心。这个中心点通常位于更平坦的盆地中,这与更好的泛化相关。平坦最小值与泛化之间的联系已被广泛研究,涉及PAC-Bayes界限和损失函数的尖锐度。
另一个视角将SWA与贝叶斯推断联系起来。在某些假设下,SGD的轨迹可以被视为从后验分布采样的马尔可夫链。沿此轨迹进行平均近似于后验均值,这是平方误差损失下的贝叶斯最优预测。这种解释将SWA与随机梯度马尔可夫链蒙特卡洛方法联系起来,尽管SWA更简单且不需要仔细调整噪声尺度。
变体与扩展
已提出多种SWA变体以进一步提高性能。一个显著的扩展是随机权重平均高斯(SWAG),由Wesley Maddox、Timur Garipov、Pavel Izmailov、Dmitry Vetrov和Andrew Gordon Wilson在2019年提出。SWAG通过使用收集权重的第一和第二矩拟合高斯分布来近似权重的后验分布。这允许不确定性估计和改进的校准,并可用于贝叶斯深度学习。
另一个变体是快速几何集成(FGE),它使用循环学习率调度从损失景观的不同模式收集权重。SWA可以与FGE结合以获得更好的性能。此外,SWA已与学习率调度、数据增强和剪枝等技术集成,以增强鲁棒性和效率。
应用
SWA已成功应用于广泛的领域。在计算机视觉中,它提高了CIFAR-10、CIFAR-100和ImageNet等数据集上图像分类器的准确性。例如,SWA已被证明可以在不改变架构的情况下,将PreAct ResNet-164在CIFAR-100上的测试准确率从81.15%提高到82.90%。在目标检测中,SWA有助于稳定训练并提高平均精度。
在自然语言处理中,SWA已被用于更有效地训练Transformer和大型语言模型。它已被证明可以减少过拟合并改善语言建模任务中的困惑度。SWA还通过提供更平滑的权重更新,使生成模型(如GAN和扩散模型)受益,这可以带来更稳定的训练和更高质量的样本。
除了标准监督学习外,SWA已被适应于半监督学习、迁移学习和领域适应。其简单性和通用性使其成为机器学习从业者工具箱中的宝贵工具。
与其他技术的比较
SWA常与其他正则化和集成方法进行比较。与平均多个独立训练模型预测的传统集成不同,SWA平均权重,这在计算上更便宜且只需要一次训练运行。这使得它在模型存储或推理成本成为关注点时特别有吸引力。
与丢弃或批归一化等技术相比,SWA不修改训练动态,而是后处理权重。它可以与这些方法结合使用。SWA也不同于Adam和其他自适应优化器,后者维护每个参数的学习率;SWA通常应用于SGD或Adam轨迹之上。
SWA的一个局限性是它需要调度来确定何时开始收集权重。如果收集开始得太早,平均值可能被次优权重主导;如果太晚,收益会减少。然而,在实践中,一个简单的经验法则是在学习率衰减到较低值后开始收集。
实际考虑
实现SWA很简单。大多数深度学习框架(如PyTorch和TensorFlow)提供SWA的实用程序或示例。主要步骤是:(1)使用标准调度训练模型,(2)继续使用恒定或循环学习率训练,(3)按间隔收集权重,(4)计算运行平均值,(5)重新计算批归一化统计量。
SWA已被证明对超参数选择具有鲁棒性。第二阶段的典型学习率设置为初始调度的最小值或一个小常数。权重收集的频率可以是每个epoch或每几百次迭代。在实践中,收集5到20个epoch的权重通常就足够了。
影响与接受度
SWA的引入对深度学习社区产生了重大影响。它提供了一种简单而有效的方法来改善泛化,而无需在推理期间增加额外计算成本。该论文被广泛引用,并启发了许多关于权重平均和平坦最小值的后续工作。SWA现在是许多训练流程中的标准技术,尤其是在学术研究和竞赛中。
SWA也影响了其他平均方法的发展,如指数移动平均(EMA),它常用于训练生成模型和大型语言模型。EMA是一种给近期权重更多权重的变体,由于其简单性和有效性,在实践中经常使用。SWA和EMA是互补的,一些框架允许在它们之间切换。
局限性与未来方向
尽管有诸多好处,SWA仍有局限性。它假设损失景观相对平滑,并且平均权重是有意义的,这可能不适用于所有架构或任务。例如,使用批归一化的模型需要仔细处理,而具有复杂参数空间的模型(如循环网络)可能受益较少。此外,除非扩展(如SWAG),SWA不提供不确定性估计。
未来的研究方向包括开发自适应平均调度、将SWA与贝叶斯方法结合,以及将SWA应用于强化学习和联邦学习等新领域。随着深度学习的不断发展,像SWA这样的权重平均技术仍然是提高模型性能和可靠性的基本工具。