随机权值平均(Stochastic Weight Averaging,SWA)
随机权值平均只需快速集合集成的一小部分算力,就可以接近其表现。SWA 可以用在任意架构和数据集上,都会有不错的表现。根据论文中的实验,SWA 可以得到我之前提到过的更宽的极小值。在经典认知下,SWA 不算集成,因为在训练的最终阶段你只得到一个模型,但它的表现超过了快照集成,接近 FGE(多个模型取平均)。
左图:W1、W2、W3分别代表3个独立训练的网络,Wswa为其平均值。中图:WSWA 在测试集上的表现超越了SGD。右图:WSWA 在训练时的损失比SGD要高。
结合 WSWA 在测试集上优于 SGD 的表现,这意味着尽管 WSWA 训练时的损失较高,它的泛化性更好。
SWA 的直觉来自以下由经验得到的观察:每个学习率周期得到的局部极小值倾向于堆积在损失平面的低损失值区域的边缘(上图左侧的图形中,褐色区域误差较低,点W1、W2、3分别表示3个独立训练的网络,位于褐色区域的边缘)。对这些点取平均值,可能得到一个宽阔的泛化解,其损失更低(上图左侧图形中的 WSWA)。
下面是 SWA 的工作原理。它只保存两个模型,而不是许多模型的集成:
- 第一个模型保存模型权值的平均值(WSWA)。在训练结束后,它将是用于预测的最终模型。
- 第二个模型(W)将穿过权值空间,基于周期性学习率规划探索权重空间。
SWA权重更新公式
在每个学习率周期的末尾,第二个模型的当前权重将用来更新第一个模型的权重(公式如上)。因此,在训练阶段,只需训练一个模型,并在内存中储存两个模型。预测时只需要平均模型,基于其进行预测将比之前描述的集成快很多,因为在那种集成中,你需要使用多个模型进行预测,最后再进行平均。
方法实现:
论文的作者自己提供了一份 PyTorch 的实现 :
此外,基于 fast.ai 库的 SWA 可见 :
Add Stochastic Weight Averaging by wdhorton · Pull Request #276 · fastai/fastaigithub.com