Test-time Batch Normalization
This work addresses distribution shift issues for deep learning models, offering a test-time optimization method that is incremental but effective in specific settings.
The paper tackles the problem of data distribution shift between training and testing in deep neural networks by proposing a novel test-time batch normalization layer called GpreBN, which optimizes during testing using entropy loss and achieves state-of-the-art results on domain generalization and robustness tasks.
Deep neural networks often suffer the data distribution shift between training and testing, and the batch statistics are observed to reflect the shift. In this paper, targeting of alleviating distribution shift in test time, we revisit the batch normalization (BN) in the training process and reveals two key insights benefiting test-time optimization: $(i)$ preserving the same gradient backpropagation form as training, and $(ii)$ using dataset-level statistics for robust optimization and inference. Based on the two insights, we propose a novel test-time BN layer design, GpreBN, which is optimized during testing by minimizing Entropy loss. We verify the effectiveness of our method on two typical settings with distribution shift, i.e., domain generalization and robustness tasks. Our GpreBN significantly improves the test-time performance and achieves the state of the art results.