[论文解读] Stochastic Weight Averaging in Parallel: Large-Batch Training that Generalizes Well

Vipul Gupta, Santiago Akle Serrano+1 more|arXiv (Cornell University)|Jan 7, 2020

Advanced Neural Network Applications参考文献 26被引用 15

一句话总结

本文提出了一种并行随机权重平均（Stochastic Weight Averaging in Parallel, SWAP）方法，通过先使用大批次快速收敛训练模型，再通过多个独立训练的小批次模型的权重平均来优化模型，从而加速深度神经网络的大批次训练。SWAP在显著缩短训练时间的同时，实现了与小批次训练相当的泛化性能，在CIFAR10上相比最先进方法将训练时间减少了高达68%。

ABSTRACT

We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm to accelerate DNN training. Our algorithm uses large mini-batches to compute an approximate solution quickly and then refines it by averaging the weights of multiple models computed independently and in parallel. The resulting models generalize equally well as those trained with small mini-batches but are produced in a substantially shorter time. We demonstrate the reduction in training time and the good generalization performance of the resulting models on the computer vision datasets CIFAR10, CIFAR100, and ImageNet.

研究动机与目标

• 解决大批次深度学习训练中训练速度与泛化性能之间的权衡问题。
• 通过在大批次初始化后并行训练多个小批次模型，实现在不牺牲模型泛化性能的前提下加快训练速度。
• 开发一种实用且低超参数调优的高效方法，以充分利用大规模DNN训练的计算资源。
• 证明大批次训练后接并行小批次微调可获得与小批次训练相当的泛化性能。

提出的方法

• 该方法首先使用大批次快速训练模型，使其迅速收敛至低训练损失区域。
• 随后，从大批次模型出发，使用小批次并行训练多个独立模型。
• 最终模型通过平均这些独立训练的小批次模型的权重得到。
• 该方法在小批次阶段采用循环学习率调度，以促进收敛至更平坦的极小值。
• 从大批次到小批次训练的过渡由一个超参数控制，实验中通过网格搜索选择。
• 该方法兼容多种优化算法，并可与局部SGD等技术结合以减少通信开销。

实验结果

研究问题

• RQ1 能否在不损害泛化性能的前提下加速大批次训练？
• RQ2 与标准大批次训练相比，对多个独立训练的小批次模型权重进行平均是否能提升泛化性能？
• RQ3 大批次预训练与并行小批次微调的结合能否实现与小批次训练相当的泛化性能？
• RQ4 所提方法的训练时间与标准小批次和大批次训练策略相比如何？
• RQ5 是否存在一个关键的过渡点，使大批次与小批次阶段的切换能同时优化性能与效率？

主要发现

• SWAP在CIFAR100上仅用241.54秒即达到79.11%的测试准确率，性能与小批次SWA相当，但耗时仅为后者的1/3.5。
• 在CIFAR10上，SWAP将训练时间减少至DAWNBench冠军参赛方法的68%，同时实现了相近的泛化性能。
• 大批次训练后接小批次SWA的训练方式，所需时间超过SWAP的三倍才能达到相似的测试准确率（CIFAR100上分别为78.12% vs. 78.18%）。
• SWAP在40个小型批次训练周期后，于241.54秒内实现CIFAR100上79.11%的测试准确率，优于大批次SWA（76.00%）和大批次后接小批次SWA（78.12%），在速度与准确率上均更优。
• SWAP的最终平均模型被发现比标准SGD模型更接近损失盆地的中心，表明其鲁棒性更强。
• 可视化结果表明，大批次模型收敛至的损失盆地与小批次模型微调后收敛的区域高度重叠，支持了该损失盆地近乎凸性的假设。