Press "Enter" to skip to content

神经网络优化中的Weight Averaging

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

 

1. 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,来使得最终的收敛位置更靠近中心的位置。在实际实验中,也显示了类似的结果,经过SWA平均之后,网络的测试准确率都有不同程度的提升。

 

2. Stochastic Weight Averaging in Parallel

 

这篇文章是对上面的SWA在并行优化中的一个应用。在并行优化神经网络的过程中,batch size的增加可以使SGD的梯度计算更精确,因此可以使用更大的lr进行优化,同时也可以缩短优化的次数。但是这样往往也会使得优化出来的generalization performance更差,因此就需要引入一些额外的技术来避免这个问题。

 

本文则是讲优化过程分成三个阶段。在前期利用large batch size的优势让网络的loss快速收敛到一个相对不错的平坦区。在第二阶段每一个节点独立的用mini batch来更新模型。最后利用SWA来对这些模型进行平均,改善large batch带来的generalization 问题。

本文在实验分析中,同样发现了类似SWA的现象就是train loss的曲面与test loss的曲面不一致。本文利用可视化方法绘制了一张CIFAR-10上的loss landscape如下图所示。

从图上可见,train loss的flat区域要比test loss大得多;同时SGD更倾向于停在flat区域的边缘。而经过SWA之后,平均之后的模型有更大的概率落在相对中心的位置上。

 

3. Lookahead Optimizer

 

前面提到的SWA在优化上,并没有改变原本优化器更新的梯度,只在结束之后选取一部分checkpoints进行weight averaging得到最终的权重。而这篇文章则是在更新过程中,利用指数移动平均的方式来计算梯度更新权重。

 

本文提出了一种权重的更新策略,每一个step的优化中维护两组权重,第一组称为fast weight,就是常规优化器更新得到的权重,第二组称为slow weight,是利用fast weight得到的权重。之所以称为fast/slow,是由于二者的更新频率不同,先用fast更新k步,然后根据得到的fast weight更新slow weight 1步作为这一个优化step的结果,依次循环进行。

在实验中,本文也通过大量的实验,验证了Lookahead优化器在前期比SGD优化的更快。此外,在实验分析中,本文也发现了一个有趣的现象,就是每一轮从slow weight开始fast weight(SGD/Adam更新)反而让loss上升了,而经过slow weight移动平均之后loss又恢复了下降。

4. Filter Grafting

 

这篇论文虽然在出发点上与前文并不相同,但是实际的方法也可以看作是一种weight average,因此也总结在这里了。

 

神经网络经常会出现冗余的问题,一般的方法都是剪枝的方式来消除冗余部分,而也有一些其他的方法则是重新利用冗余部分来提高网络的性能/泛化性,例如Dense sparse dense training、RePr等等。本文也是同样的出发点,希望通过引入外部信息来改善冗余问题。

 

本文提出了一种利用权重的熵来评估网络中filter所包含的信息量,在优化中,同时优化两个相同的网络,采用不同的超参来进行优化,在优化过程中对filter进行加权平均实现对信息量的补足。

在加权平均时,为了简洁,不采用针对特定位置的filter进行加权,而是根据一层的熵大小来对整个一层的参数进行平均,加权所用的alpha则是根据两个网络这一层熵的大小自适应决定的。更进一步,则是可以将2个网络相互加权,拓展成多个网络循环加权,如下图所示。

在实验中,本文除了对这种grafting策略进行了性能测试,也对其他的一些细节进行了分析:(1) 不同的信息来源对提升的影响(网络自身、噪声、不同网络);(2) 不同的信息量评估方式的影响(L1 norm、熵)。最终得出文中提出的多个网络基于熵的杂交策略是最优的。同时也对杂交训练得到的网络的冗余量(权重L1 norm统计)和网络最终熵之和进行了分析。

[1] Izmailov P, Podoprikhin D, Garipov T, et al. Averaging weights leads to wider optima and better generalization[J]. arXiv preprint arXiv:1803.05407, 2018.

 

[2] Gupta V, Serrano S A, DeCoste D. Stochastic Weight Averaging in Parallel: Large-Batch Training that Generalizes Well[J]. arXiv preprint arXiv:2001.02312, 2020.

 

[3] Zhang M, Lucas J, Ba J, et al. Lookahead optimizer: k steps forward, 1 step back[J]. Advances in Neural Information Processing Systems, 2019, 32: 9597-9608.

 

[4] Meng F, Cheng H, Li K, et al. Filter grafting for deep neural networks[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020: 6599-6607.

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注