神经网络优化中的Weight Averaging

©PaperWeekly 原创 · 作者｜张子逊

研究方向｜神经网络剪枝、NAS

在神经网络优化的研究中，有研究改进优化器本身的（例如学习率衰减策略、一系列 Adam 改进等等），也有不少是改进 normalization 之类的技术（例如 Weight Decay、BN、GN 等等）来提高优化器的性能和稳定性。除此之外，还有一个比较常见的技术就是 Weight Averaging，也就是字面意思对网络的权重进行平均，这也是一个不错的提高优化器性能/稳定性的方式。

Stochastic Weight Averaging (SWA)

在神经网络的优化中，有一个比较公认的问题就是 train loss 和 test loss 优化曲面不一致。导致这个问题可能的原因有很多，（以下是个人推测）可能是 train 和 test 数据本身存在分布不一致的问题、也可能是因为训练的时候在 train loss 中加入了一系列正则化等等。由于这个不一致的问题，就会导致优化出来的网络 generalization performance 可能会不好。 

本文提出了一个比较简单直接的方式来解决这个问题，在优化的末期取 k 个优化轨迹上的 checkpoints，平均他们的权重，得到最终的网络权重，这样就会使得最终的权重位于 flat 曲面更中心的位置。这个方法也被集成到了 PyTorch 1.6 之后的版本中。

本文的实验分析部分也给出了详细的分析和结果来验证这样的一种方法是有效的。首先是分析的部分，本文通过可视化绘制了通过 SWA 和 SGD 在 loss landscape 上的收敛位置，以及 SGD 优化得到的权重对应的 loss 和 error 相比 SWA 得到的权重的距离，如下图所示。

从图上可以看到几个有趣的现象：首先，train loss 和 test loss 的 landscape 之间确实存在偏移；其次，SGD 更倾向于收敛到 flat 区域的边缘。比较直观的一种猜想就是，利用这样的性质，SWA 可以通过平均 flat 区域边缘的一些 checkpoints，来使得最终的收敛位置更靠近中心的位置。在实际实验中，也显示了类