Admissibility of Stein Shrinkage for Batch Normalization in the Presence of Adversarial Attacks

TL;DR

This paper demonstrates that Stein shrinkage estimators improve batch normalization's robustness and accuracy, especially under adversarial attacks, leading to state-of-the-art results in image classification and segmentation tasks.

Contribution

It proves the dominance of Stein shrinkage estimators over traditional methods for BN parameters under adversarial conditions and shows improved robustness and performance in deep learning models.

Findings

Stein shrinkage estimators outperform sample mean and variance in adversarial settings.

James-Stein BN yields smaller Lipschitz constants, indicating better regularity.

Stein-corrected BN achieves SOTA results on multiple datasets.

Abstract

Batch normalization (BN) is a ubiquitous operation in deep neural networks, primarily used to improve stability and regularization during training. BN centers and scales feature maps using sample means and variances, which are naturally suited for Stein's shrinkage estimation. Applying such shrinkage yields more accurate mean and variance estimates of the batch in the mean-squared-error sense. In this paper, we prove that the Stein shrinkage estimator of the mean and variance dominates over the sample mean and variance estimators, respectively, in the presence of adversarial attacks modeled using sub-Gaussian distributions. Furthermore, by construction, the James-Stein (JS) BN yields a smaller local Lipschitz constant compared to the vanilla BN, implying better regularity properties and potentially improved robustness. This facilitates and justifies the application of Stein shrinkage to estimate the mean and variance parameters in BN and the use of it in image classification and segmentation tasks with and without adversarial attacks. We present SOTA performance results using this Stein-corrected BN in a standard ResNet architecture applied to the task of image classification using CIFAR-10 data, 3D CNN on PPMI (neuroimaging) data, and image segmentation using HRNet on Cityscape data with and without adversarial attacks.