[论文解读] Switch EMA: A Free Lunch for Better Flatness and Sharpness

Siyuan Li, Zicheng Liu+10 more|arXiv (Cornell University)|Feb 14, 2024

Soil Mechanics and Vehicle DynamicsEngineering被引用 3

一句话总结

本文提出Switch EMA（SEMA），一种简单而高效的方法，通过在每个训练周期结束时动态切换指数移动平均（EMA）参数到主模型，实现性能提升。通过在快速收敛的锐利更新与平坦、泛化性能优化的EMA状态之间交替，SEMA在视觉和自然语言处理任务中实现了更优的模型泛化能力和更快的收敛速度，且无需额外计算成本——为深度学习训练提供了一种即插即用的‘免费午餐’。

ABSTRACT

Exponential Moving Average (EMA) is a widely used weight averaging (WA) regularization to learn flat optima for better generalizations without extra cost in deep neural network (DNN) optimization. Despite achieving better flatness, existing WA methods might fall into worse final performances or require extra test-time computations. This work unveils the full potential of EMA with a single line of modification, i.e., switching the EMA parameters to the original model after each epoch, dubbed as Switch EMA (SEMA). From both theoretical and empirical aspects, we demonstrate that SEMA can help DNNs to reach generalization optima that better trade-off between flatness and sharpness. To verify the effectiveness of SEMA, we conduct comparison experiments with discriminative, generative, and regression tasks on vision and language datasets, including image classification, self-supervised learning, object detection and segmentation, image generation, video prediction, attribute regression, and language modeling. Comprehensive results with popular optimizers and networks show that SEMA is a free lunch for DNN training by improving performances and boosting convergence speeds.

研究动机与目标

• 解决现有权重平均（WA）方法的局限性：虽然能提升损失曲面的平坦度，但常无法提升最终性能，或需要额外推理开销。
• 统一快速收敛（锐利度）与泛化能力（平坦度）的优势，形成一种高效训练策略。
• 提出一种极简修改：在每个训练周期结束后，将EMA参数切换至主模型，从而同时增强优化动态与最终模型性能。
• 证明SEMA可作为‘免费午餐’，在不增加计算成本的前提下，显著提升各类任务与架构下的准确率与收敛速度。

提出的方法

• SEMA引入一种动态切换机制：在每个训练周期结束时，将EMA平均后的权重复制到主模型的参数中。
• 该机制形成双阶段训练循环：主模型快速探索损失曲面的锐利区域，而EMA模型则追踪一条更平坦、更具泛化能力的路径。
• 切换过程确保每个新周期均从平坦、泛化优化的状态开始，从而实现更快地收敛至更优的极小值。
• 该方法完全即插即用，仅需对现有EMA实现进行一行代码修改即可集成。
• SEMA完全在训练阶段运行，避免了SWA或基于集成方法中常见的推理时开销。
• 该方法充分利用了快速锐利更新与稳定平坦EMA状态的互补优势，实现优化动态的平衡。

实验结果

研究问题

• RQ1 在深度神经网络优化中，仅通过在EMA与主模型权重之间简单切换，能否同时改善平坦度与锐利度？
• RQ2 在每个周期结束后将EMA参数切换至主模型，是否能实现比标准EMA或基线训练更好的泛化能力与更快的收敛速度？
• RQ3 SEMA是否能在不引入额外计算成本或推理开销的前提下实现性能提升？
• RQ4 在多样化任务中，SEMA与SWA、SAM或EMA等现有方法相比，在优化动态与最终模型准确率方面表现如何？

主要发现

• SEMA在12个多样化基准测试中持续提升测试准确率，涵盖图像分类、目标检测、自监督学习、视频预测与语言建模任务。
• 在DeiT-S模型的ImageNet-1K数据集上，SEMA达到84.2%的top-1准确率，优于基线（83.8%）与EMA（84.0%），且收敛更快。
• 在ResNet-50 Cascade Mask R-CNN的COCO目标检测任务中，SEMA相较基线提升AP 0.7%，相较EMA提升0.5%。
• 在MoCo.V3的CIFAR-100上，SEMA将线性探测准确率提升至88.9%，超越EMA（88.3%）与基线（87.6%）。
• 在AgeDB数据集的年龄回归任务中，SEMA相较EMA降低MAE 0.4年，相较基线降低0.6年。
• 在所有评估任务中，SEMA均显著加速收敛速度，训练曲线显示其性能提升更早且更陡峭，优于EMA与基线模型。