[论文解读] Towards Stabilizing Batch Statistics in Backward Propagation of Batch Normalization

Junjie Yan, Ruosi Wan+4 more|arXiv (Cornell University)|Jan 19, 2020

Blind Source Separation Techniques参考文献 17被引用 29

一句话总结

本文提出移动平均批归一化（MABN），一种新型归一化方法，通过在反向传播中用移动平均值替代批统计量，稳定了批归一化过程，使原始批归一化在小批量设置下的性能得以完全恢复，且推理阶段不引入非线性操作。MABN 在批量大小为 1 或 2 时仍能达到与标准 BN 相当的性能，在 COCO 和 ImageNet 上与 SyncBN 表现相当，同时保持了高效的推理速度。

ABSTRACT

Batch Normalization (BN) is one of the most widely used techniques in Deep Learning field. But its performance can awfully degrade with insufficient batch size. This weakness limits the usage of BN on many computer vision tasks like detection or segmentation, where batch size is usually small due to the constraint of memory consumption. Therefore many modified normalization techniques have been proposed, which either fail to restore the performance of BN completely, or have to introduce additional nonlinear operations in inference procedure and increase huge consumption. In this paper, we reveal that there are two extra batch statistics involved in backward propagation of BN, on which has never been well discussed before. The extra batch statistics associated with gradients also can severely affect the training of deep neural network. Based on our analysis, we propose a novel normalization method, named Moving Average Batch Normalization (MABN). MABN can completely restore the performance of vanilla BN in small batch cases, without introducing any additional nonlinear operations in inference procedure. We prove the benefits of MABN by both theoretical analysis and experiments. Our experiments demonstrate the effectiveness of MABN in multiple computer vision tasks including ImageNet and COCO. The code has been released in https://github.com/megvii-model/MABN.

研究动机与目标

• 为解决批归一化（BN）在小批量训练中出现的不稳定性问题，尤其是在目标检测与实例分割等受显存限制而批量大小受限的视觉任务中。
• 识别并分析 BN 反向传播中此前被忽视的两类批统计量，揭示其对训练不稳定性的贡献。
• 开发一种归一化方法，在不引入推理阶段非线性操作的前提下，恢复原始 BN 在小批量场景下的完整性能。
• 在多个视觉基准上提供理论依据与实证验证，支持所提方法的有效性。

提出的方法

• 该方法引入两种独立的移动平均统计量：前向传播中使用指数移动平均统计量（EMAS），反向传播中使用平滑移动平均统计量（SMAS）。
• EMAS 在前向传播中替代批统计量，通过动量更新运行统计量，以稳定特征归一化。
• SMAS 在反向传播中替代批统计量，以稳定梯度计算，并采用改进的归一化形式以防止训练崩溃。
• 重新组织归一化形式，减少批统计量数量，并集中卷积核权重，以提升稳定性。
• 采用重归一化策略以保持梯度一致性，防止小批量训练中出现发散。
• 使用基于动量的更新规则处理移动统计量，超参数经调优以确保稳定性和收敛性。

实验结果

研究问题

• RQ1 为何批归一化在小批量训练中会失效，尽管其在大批量设置下表现优异？
• RQ2 批归一化反向传播中被忽视的批统计量具体是什么？它们如何影响训练稳定性？
• RQ3 能否在不引入推理阶段非线性操作的前提下，完全恢复原始批归一化在小批量场景下的性能？
• RQ4 是否存在一种方法，通过在前向与反向传播中均使用移动平均值来稳定统计量，同时保持 BN 的线性特性以实现高效推理？

主要发现

• 在使用 ResNet-50 的 ImageNet 上，MABN 在批量大小为 2 时达到 23.58% 的 Top-1 验证误差，与大批次下标准 BN 的性能完全一致。
• 在 COCO 实例分割任务中，MABN 达到 AP^{bbox} 为 34.85 和 AP^{mask} 为 31.61，与 SyncBN 表现相当，显著优于原始 BN 和 BRN。
• 消融实验表明，前向传播中的 EMAS 与反向传播中的 SMAS 必须同时存在，移除任一组件均导致性能显著下降。
• MABN 的推理速度与原始 BN 完全一致，而实例归一化方法因引入非线性操作导致推理时间翻倍。
• 使用 MABN 训练时，即使在批量大小为 1 的极端情况下也能实现稳定收敛，而直接在反向传播中使用移动平均值则因梯度不稳定性而失败。
• 理论分析表明，MABN 所采用的改进归一化形式在小批量设置下比原始 BN 形式更具稳定性。