[Paper Review] Evaluating Prediction-Time Batch Normalization for Robustness under Covariate Shift

Zachary Nado, Shreyas Padhy+4 more|arXiv (Cornell University)|Jun 19, 2020

Domain Adaptation and Few-Shot LearningComputer Science52 references95 citations

TL;DR

The paper introduces prediction-time batch normalization to recalibrate activations using a small unlabeled batch at prediction time, improving accuracy and calibration under covariate shift, with strong results on CIFAR-10-C and ImageNet-C.

ABSTRACT

Covariate shift has been shown to sharply degrade both predictive accuracy and the calibration of uncertainty estimates for deep learning models. This is worrying, because covariate shift is prevalent in a wide range of real world deployment settings. However, in this paper, we note that frequently there exists the potential to access small unlabeled batches of the shifted data just before prediction time. This interesting observation enables a simple but surprisingly effective method which we call prediction-time batch normalization, which significantly improves model accuracy and calibration under covariate shift. Using this one line code change, we achieve state-of-the-art on recent covariate shift benchmarks and an mCE of 60.28\% on the challenging ImageNet-C dataset; to our knowledge, this is the best result for any model that does not incorporate additional data augmentation or modification of the training pipeline. We show that prediction-time batch normalization provides complementary benefits to existing state-of-the-art approaches for improving robustness (e.g. deep ensembles) and combining the two further improves performance. Our findings are supported by detailed measurements of the effect of this strategy on model behavior across rigorous ablations on various dataset modalities. However, the method has mixed results when used alongside pre-training, and does not seem to perform as well under more natural types of dataset shift, and is therefore worthy of additional study. We include links to the data in our figures to improve reproducibility, including a Python notebooks that can be run to easily modify our analysis at https://colab.research.google.com/drive/11N0wDZnMQQuLrRwRoumDCrhSaIhkqjof.

Motivation & Objective

• Motivate and formalize the prediction-time batch setting where predictions occur in small batches at test time under covariate shift.
• Propose a simple, efficient method—prediction-time BN—to recalibrate activations using the current prediction batch statistics.
• Evaluate the method across image and non-image modalities on covariate shift benchmarks and analyze when it helps or fails.

Proposed method

• Formalize the prediction-time batch setting with batch-wise loss and risk minimization.
• Apply batch normalization statistics recalculated on each prediction-time batch (prediction-time BN) as opposed to frozen training EMA statistics.
• Compare prediction-time BN with vanilla BN, ensembles, temperature scaling, and other normalization variants across multiple datasets.
• Provide ablations to understand the role of epsilon, which BN layers to reset, and interaction with pretraining and natural shifts.

Experimental results

Research questions

• RQ1 Does recalculating batch norm statistics on prediction-time batches improve calibration and accuracy under covariate shift?
• RQ2 How does prediction-time BN compare to training-time BN and other calibration methods across image and non-image modalities?
• RQ3 What are the limitations and failure modes of prediction-time BN, including effects of pretraining and natural shifts?
• RQ4 How sensitive is the method to batch size, choice of BN layer, and normalization hyperparameters?

Key findings

• Prediction-time BN aligns activation supports of shifted data with training statistics, improving calibration and often accuracy under covariate shift.
• On CIFAR-10-C and ImageNet-C, prediction-time BN yields strong calibration and competitive accuracy, with an mCE of 60.28% on ImageNet-C without extra data augmentation.
• The method is complementary to ensembles and maintains gains across a range of prediction batch sizes, with substantial benefits even for modest batch sizes (~100).
• Prediction-time BN may underperform when used with pretraining (e.g., Noisy Student on ImageNet-C) and under more natural dataset shifts, indicating boundary conditions for effectiveness.
• In natural adversarial datasets (ImageNet-A), prediction-time BN improves calibration and can outperform train BN in some settings.
• Ablation studies show normalization layers before the output are insufficient alone; re-normalizing internal BN layers yields better gains.