Towards Stabilizing Batch Statistics in Backward Propagation of Batch Normalization

TL;DR

This paper identifies overlooked batch statistics in BN's backward pass that impair training with small batches and introduces MABN, a method restoring BN's performance without extra inference complexity, validated on vision tasks.

Contribution

The paper reveals two previously unconsidered batch statistics in BN's backward propagation and proposes MABN, a novel normalization method that maintains BN's performance with small batch sizes.

Findings

MABN restores BN performance in small batch scenarios.

Theoretical analysis supports MABN's effectiveness.

Experiments on ImageNet and COCO validate improvements.

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.