PaperlayerView exact revision on arXiv

Source-linked AI summary

Evaluating Prediction-Time Batch Normalization for Robustness under Covariate Shift

Zachary Nado, Shreyas Padhy, D. Sculley, Alexander D'Amour, Balaji Lakshminarayanan, Jasper Snoek

arXiv:2006.10963v3cs.LGstat.ML

TL;DR

Covariate shift can sharply degrade deep models' accuracy and uncertainty calibration, motivating methods that use accessible unlabeled shifted data at prediction time. The paper proposes prediction-time batch normalization, which recomputes normalization statistics for each test batch and reports strong robustness results, including 60.28% mCE on ImageNet-C. Its benefits are practical but bounded by weaker performance on natural shifts and alongside pre-training.

  • Problem

    Covariate shift degrades deep models' predictive accuracy and calibration, while practical deployments often encounter shifted data at prediction time.

  • Method

    Prediction-time batch normalization recomputes batch-normalization statistics for each unlabeled test batch instead of using frozen training statistics.

  • Results

    The method works effectively across multiple modalities and achieves 60.28% mCE on ImageNet-C without additional data augmentation.

  • Takeaways & Limitations

    Prediction-time batch normalization is a simple, computationally efficient approach for correcting covariate shift in real-world deployment settings, with caveats.

  • Takeaways & Limitations

    Performance is weaker on more natural dataset shifts and decreases when prediction-time batch normalization is combined with pre-training.

Abstract

from arXiv · show

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.

1 Introduction

Covariate shift undermines deep models' predictions and confidence, but prediction-time batches provide a brief opportunity for efficient, unlabeled correction without retraining. The paper proposes prediction-time batch normalization and reports strong performance across practical shift settings.

  • Covariate shift breaks i.i.d. assumptions and can produce incorrect, disproportionately overconfident predictions in deep models.
  • Prediction-time batching exposes small unlabeled samples whose distribution can reflect local deployment conditions.Modern systems batch examples for computational efficiency, creating a short window to inspect shifted data characteristics.
  • Practical correction must be computationally efficient and avoid retraining because prediction-time access lasts only milliseconds.
  • Prediction-time batch normalization extends training-time batch normalization to correct covariate shift using prediction-time batch information.The paper motivates the method by analyzing hidden-layer activations under shift and evaluates it across multiple modalities.
  • 60.28% mCE is achieved on ImageNet-C without additional data augmentation, reported as the best result to the authors' knowledge under that constraint.

2 Setup

The paper formalizes prediction as batch-wise inference on an unknown target distribution that may differ from training in its feature marginal. It evaluates predictive risk and uncertainty quality using standard batch losses, calibration error, and Brier score.

  • Training examples come from p(x, y), while unlabeled test examples come from a potentially distinct target distribution q(x) with q(y | x).
  • Prediction-time inference operates on batches of t < T examples, and the model produces their predictions simultaneously.
  • The objective is expected batch-wise loss over i.i.d. test batches, with the batch loss decomposing additively across examples.Because the loss is linear, empirical evaluation uses the mean per-example loss in standard pipelines.
  • Covariate shift is defined by p(x) ≠ q(x) with unchanged conditional labels p(y | x) = q(y | x).Effective uncertainty should be lower for inputs rare under or outside the training feature distribution.
  • Expected Calibration Error measures confidence-accuracy gaps across confidence bins, while Brier Score measures squared distance from one-hot targets.Brier Score decreases toward zero as the predictive distribution approaches the true distribution.

3 Related work

Prior work addresses covariate shift through domain adaptation, batch-normalization statistics, calibration strategies, and prediction-time batch methods, but this paper focuses on test-time intervention for discriminative models.

  • Domain-adaptation methods commonly assume access to unlabeled target samples during training, whereas this paper intervenes after model training.
  • AdaBN computes domain-specific batch-normalization statistics using the target domain at test time, although related methods often require target data during training.
  • Existing calibration improvements for image benchmarks often rely on extensive data augmentation or substantial pre-training.
  • The paper notes a conjectured worst-case generalization dependence on the number of batches rather than the number of test examples.
  • Prediction-time batch normalization has been explored for generative-model OOD detection, but this paper studies discriminative models under covariate shift.

4 Prediction-Time Batch Normalization

Prediction-time batch normalization recalculates normalization statistics on each unlabeled prediction batch to correct shifted activations. The method aligns activation supports and is associated with improved uncertainty quantification and lower Brier scores under covariate shift.

  • Prediction-time BN recomputes batch-normalization statistics for each test batch instead of using frozen training-time statistics.The normalization operates on pre-activations within each batch.
  • The experiments find prediction-time BN surprisingly effective for improving uncertainty quantification in deep networks.
  • Covariate shift can move internal activations outside the ranges encountered during training, worsening calibration as activation-distance from training increases.The paper measures this behavior using Brier Score and relates it to shifted hidden-layer activations.
  • Prediction-time BN brings shifted penultimate-layer activation supports into alignment relative to train BN and other normalization schemes.In Figure 2, corrected supports cluster around 0 on the horizontal discrepancy axis.
  • Prediction-time BN maps out-of-support activations toward training-support regions associated with uncertain predictions, producing consistently lower Brier scores.Aligning supports alone is not sufficient; the associated prediction regions also matter for calibration.

5 Performance Under Covariate Shift

The paper evaluates prediction-time batch normalization across image and tabular covariate-shift benchmarks, comparing it with other normalization, calibration, and robustness methods. It improves performance in several settings, but shows weaker or mixed results for pre-trained models and more natural dataset shifts.

  • Evaluation setup: Prediction-time BN is evaluated on CIFAR-10-C, ImageNet-C, and corrupted Criteo data using test-time-only modifications.The image benchmarks include 19 corruption types at five shift intensities, while Criteo provides a second modality.
  • Main results: 60.28% mCE is achieved on ImageNet-C with prediction-time BN and a vanilla ResNet-50 model.Prediction-time BN is also reported as complementary to other robustness methods.
  • Comparisons: Prediction-time BN is complementary to ensembles, while temperature scaling and alternative normalization schemes provide comparison baselines.The evaluated normalization methods include batch, instance, layer, and group normalization; ensembles average predictions from M = 10 independently trained models.
  • Prediction batch dependence: Strong performance is obtained with a prediction batch of 100 examples, after which larger batches yield only marginal improvements.This addresses the concern that predictions depend on other examples and batch size.
  • Prediction batch dependence: Frozen prediction-time BN retains substantial improvements using statistics from a single batch per shifted split, although simultaneous shift types reduce its relative advantage.Across up to 19 simultaneous corruption types, prediction-time BN still outperforms train BN.
  • Limitations: Prediction-time BN performs worse than train BN for the pre-trained Noisy Student EfficientNet on ImageNet-C.The paper suggests that extensive pre-training may expose train-time statistics to patterns resembling the benchmark corruptions.
  • Limitations: On ImageNet-v2, prediction-time statistics slightly reduce accuracy and Brier Score but improve ECE; on ImageNet-A, prediction-time BN outperforms train BN.These results indicate different behavior across more natural or adversarial dataset shifts.

6 Ablation Studies

The ablations examine normalization choices, epsilon sensitivity, confidence behavior, activation alignment, and the role of internal versus final normalization layers under covariate shift.

  • Sensitivity to ϵ: Higher prediction-time epsilon values improve calibration, but increasing epsilon beyond roughly two orders of magnitude causes accuracy to collapse.The denominator variance term acts like an adaptive, unsupervised temperature that changes per batch.
  • Comparing Normalization Methods: Prediction-time BN improves calibration and accuracy under CIFAR-10-C shift, whereas InstanceNorm improves calibration at the expense of accuracy.
  • Activation Alignment: Prediction-time BN aligns hidden and output activation distributions, but matching support does not fully match distribution densities.The remaining lower confidence on ImageNet-C is consistent with imperfect density alignment.
  • Confidence Behavior: Prediction-time BN produces slightly lower confidence and higher per-bin accuracies on ImageNet-C, avoiding excessive underconfidence.
  • Batch Normalization Architectures: Using prediction-time statistics only in the final normalization layer degrades performance, showing that preceding hidden-layer misalignment accumulates.
  • Batch Normalization Architectures: Removing all normalization layers except one before the final linear layer recovers most, but not all, performance benefits under covariate shift.Normalizing the final linear-layer inputs still significantly improves performance.

7 Conclusion

The paper attributes miscalibration under covariate shift to shifted internal activations and proposes prediction-time batch normalization to normalize them using small unlabeled prediction batches.

  • The paper hypothesizes that covariate shift moves internal activations outside the ranges encountered during training.
  • Prediction-time batch normalization is presented as the most effective normalization approach for bringing shifted activations into ranges expected by the model.
  • The method often outperforms training-statistic EMA on image and categorical data, but performs worse on more natural shifts and with pre-training.
  • Prediction-time BN requires access to a reasonably sized batch of unlabeled data immediately before prediction.

A.1 CIFAR-10

The CIFAR-10 experiments use standard augmentation during training and evaluate CIFAR-10-C without preprocessing across standard and additional corruption types.

  • CIFAR-10 training uses zero-padding, random cropping, random horizontal flips, and rescaling images to [-1, 1].
  • CIFAR-10-C receives no preprocessing and includes 15 standard corruption types plus gaussian blur, saturate, spatter, and speckle noise.

A.2 ImageNet

The ImageNet experiments use 224x224 images with standard Inception augmentation for training, while ImageNet-C is evaluated without preprocessing on all 19 corruption types.

  • ImageNet training uses 224x224 images with standard Inception data augmentation.
  • ImageNet-C receives no preprocessing and includes all 19 corruption types.

A.3 Criteo

The Criteo experiments simulate covariate shift by randomizing categorical features, using a model architecture based on the prior work of Ovadia et al. [2019].

  • A.3 Criteo: Criteo covariate shift is simulated by randomizing features with probabilities from 5% to 95%.These probabilities correspond to the shifts shown in Figure A4.
  • A.3 Criteo: The Criteo model encodes each categorical feature into a dense vector and concatenates the resulting representations.The concatenated feature vector is then processed by batch normalization and fully connected layers.
  • A.3 Criteo: The network uses three fully connected layers with widths [2572, 1454, 1596] and ReLU nonlinearities.

C Hyperparameter Tuning Ranges

The supplied material documents hyperparameter-search procedures, training configurations, normalization variants, and extensive calibration and accuracy ablations across CIFAR-10-C and ImageNet-C.

  • C Hyperparameter Tuning Ranges: 100 random trials tune the learning rate α, one minus the momentum 1 −γ, and Adam’s ϵ on logarithmic scales.The search ranges are referenced in Table 1.
  • Training configurations: CIFAR-10 training uses 100 epochs with batch size 512 and learning-rate reductions at epochs 40, 60, 80, and 90.
  • Training configurations: ImageNet training uses 90 epochs with batch size 512, five warm-up epochs, and learning-rate reductions at epochs 30, 60, and 80.
  • Training configurations: Criteo training uses one epoch with batch size 1024 and the CIFAR-10 learning-rate schedule.
  • Evaluation and ablations: Prediction-time batch normalization is evaluated against vanilla batch normalization and temperature scaling across calibration and accuracy experiments under covariate shift.Experiments cover CIFAR-10-C and ImageNet-C with test batch sizes of 500 and 100, respectively.
  • Evaluation and ablations: Performance degrades linearly as train and test activation distributions become farther apart, with higher shift intensities associated with greater activation-support distance.
  • Evaluation and ablations: On ImageNet-C, prediction-time batch normalization with ensembles achieves the best accuracy and Brier Score, while degrading ECE.The figure caption attributes this discrepancy likely to the number of confidence bins used for ECE.
  • Evaluation and ablations: Across normalization and batch-size ablations, FixUp and train batch normalization degrade as shift increases, while ImageNet-C performance plateaus after batch size 100 or 250.Accuracy and Brier Score continue to improve slightly after batch size 100 on ImageNet-C, despite a small ECE increase.
Loading 2006.10963v3…